├── LICENSE ├── MACE_models ├── mace01.model ├── mace01_compiled.model ├── mace01_run-123.log ├── mace01_run-123_train.txt ├── mace01_swa.model ├── mace01_swa_compiled.model ├── mace02_com1.model ├── mace02_com1_compiled.model ├── mace02_com1_gen1.model ├── mace02_com1_gen1_compiled.model ├── mace02_com1_gen1_run-123.log ├── mace02_com1_gen1_run-123_train.txt ├── mace02_com1_gen1_swa.model ├── mace02_com1_gen1_swa_compiled.model ├── mace02_com1_run-123.log ├── mace02_com1_run-123_train.txt ├── mace02_com1_swa.model ├── mace02_com1_swa_compiled.model ├── mace02_com2.model ├── mace02_com2_compiled.model ├── mace02_com2_run-345.log ├── mace02_com2_run-345_train.txt ├── mace02_com2_swa.model ├── mace02_com2_swa_compiled.model ├── mace02_com3.model ├── mace02_com3_compiled.model ├── mace02_com3_run-567.log ├── mace02_com3_run-567_train.txt ├── mace02_com3_swa.model └── mace02_com3_swa_compiled.model ├── README.md ├── T01-MACE-Practice-I.ipynb ├── T02-MACE-Practice-II.ipynb ├── T03-MACE-Theory.ipynb ├── T04-MLIP-Apps.ipynb ├── config ├── config-01.yml └── config-02.yml ├── data.py ├── data ├── LiFePO4_supercell.cif ├── mace02_md_100_xtb.xyz ├── solvent_configs.xyz ├── solvent_liquid.xyz ├── solvent_molecs.xyz ├── solvent_rotated.xyz ├── solvent_xtb.xyz ├── solvent_xtb_test.xyz ├── solvent_xtb_train_20.xyz ├── solvent_xtb_train_200.xyz ├── solvent_xtb_train_50.xyz └── solvent_xtb_train_53_gen1.xyz ├── figures ├── active_learning.png ├── embedding_layer.png ├── interaction_first_layer.png ├── interaction_layer.png ├── interaction_layer_second.png ├── iterative_training.png ├── mace_architecture.png ├── mace_architecture_layers.png ├── product_layer.png ├── readout.png └── sketch.png ├── moldyn ├── mace01_md.xyz ├── mace01_md_liquid.xyz ├── mace02_md.xyz ├── mace02_md_committee.xyz ├── mace02_md_gen1.xyz ├── mace03_md.xyz └── xtb_md.xyz └── tests └── mace01 ├── solvent_test.xyz └── solvent_train.xyz /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Ilyes Batatia, Will Baldwin, Ioan-Bogdan Magdau, Alin Elena 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /MACE_models/mace01.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace01.model -------------------------------------------------------------------------------- /MACE_models/mace01_compiled.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace01_compiled.model -------------------------------------------------------------------------------- /MACE_models/mace01_run-123.log: -------------------------------------------------------------------------------- 1 | 2024-07-21 14:16:39.186 INFO: MACE version: 0.3.6 2 | 2024-07-21 14:16:39.187 INFO: Configuration: Namespace(config=None, name='mace01', seed=123, log_dir='MACE_models', model_dir='MACE_models', checkpoints_dir='MACE_models', results_dir='MACE_models', downloads_dir='downloads', device='cuda', default_dtype='float64', distributed=False, log_level='INFO', error_table='PerAtomRMSE', model='MACE', r_max=4.0, radial_type='bessel', num_radial_basis=8, num_cutoff_basis=5, pair_repulsion=False, distance_transform='None', interaction='RealAgnosticResidualInteractionBlock', interaction_first='RealAgnosticResidualInteractionBlock', max_ell=2, correlation=2, num_interactions=2, MLP_irreps='16x0e', radial_MLP='[64, 64, 64]', hidden_irreps='128x0e + 128x1o', num_channels=32, max_L=0, gate='silu', scaling='rms_forces_scaling', avg_num_neighbors=1, compute_avg_num_neighbors=True, compute_stress=False, compute_forces=True, train_file='data/solvent_xtb_train_200.xyz', valid_file=None, valid_fraction=0.1, test_file='data/solvent_xtb_test.xyz', test_dir=None, multi_processed_test=False, num_workers=0, pin_memory=True, atomic_numbers=None, mean=None, std=None, statistics_file=None, E0s='average', keep_isolated_atoms=False, energy_key='energy_xtb', forces_key='forces_xtb', virials_key='virials', stress_key='stress', dipole_key='dipole', charges_key='charges', loss='weighted', forces_weight=100.0, swa_forces_weight=100.0, energy_weight=1.0, swa_energy_weight=1000.0, virials_weight=1.0, swa_virials_weight=10.0, stress_weight=1.0, swa_stress_weight=10.0, dipole_weight=1.0, swa_dipole_weight=1.0, config_type_weights='{"Default":1.0}', huber_delta=0.01, optimizer='adam', beta=0.9, batch_size=10, valid_batch_size=10, lr=0.01, swa_lr=0.001, weight_decay=5e-07, amsgrad=True, scheduler='ReduceLROnPlateau', lr_factor=0.8, scheduler_patience=50, lr_scheduler_gamma=0.9993, swa=True, start_swa=None, ema=False, ema_decay=0.99, max_num_epochs=50, patience=2048, foundation_model=None, foundation_model_readout=True, eval_interval=2, keep_checkpoints=False, save_all_checkpoints=False, restart_latest=False, save_cpu=False, clip_grad=10.0, wandb=False, wandb_dir=None, wandb_project='', wandb_entity='', wandb_name='', wandb_log_hypers=['num_channels', 'max_L', 'correlation', 'lr', 'swa_lr', 'weight_decay', 'batch_size', 'max_num_epochs', 'start_swa', 'energy_weight', 'forces_weight']) 3 | 2024-07-21 14:16:39.220 INFO: CUDA version: 12.1, CUDA device: 0 4 | 2024-07-21 14:16:39.276 INFO: Current Git commit: 30003ddde736a44f977bccdc497a4b3432e38b03 5 | 2024-07-21 14:16:39.368 INFO: Since ASE version 3.23.0b1, using stress_key 'stress' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_stress'. You need to use --stress_key='REF_stress', to tell the key name chosen. 6 | 2024-07-21 14:16:39.369 INFO: Using isolated atom energies from training file 7 | 2024-07-21 14:16:39.389 INFO: Loaded 200 training configurations from 'data/solvent_xtb_train_200.xyz' 8 | 2024-07-21 14:16:39.389 INFO: Using random 10.0% of training set for validation 9 | 2024-07-21 14:16:39.723 INFO: Since ASE version 3.23.0b1, using stress_key 'stress' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_stress'. You need to use --stress_key='REF_stress', to tell the key name chosen. 10 | 2024-07-21 14:16:39.755 INFO: Loaded 1000 test configurations from 'data/solvent_xtb_test.xyz' 11 | 2024-07-21 14:16:39.755 INFO: Total number of configurations: train=180, valid=20, tests=[Default: 1000] 12 | 2024-07-21 14:16:39.756 INFO: AtomicNumberTable: (1, 6, 8) 13 | 2024-07-21 14:16:39.757 INFO: Atomic energies: [-10.707211383396714, -48.847445262804705, -102.57117256025786] 14 | 2024-07-21 14:16:39.920 INFO: WeightedEnergyForcesLoss(energy_weight=1.000, forces_weight=100.000) 15 | 2024-07-21 14:16:40.058 INFO: Average number of neighbors: 9.86205556634933 16 | 2024-07-21 14:16:40.058 INFO: Selected the following outputs: {'energy': True, 'forces': True, 'virials': False, 'stress': False, 'dipoles': False} 17 | 2024-07-21 14:16:40.114 INFO: Building model 18 | 2024-07-21 14:16:40.115 INFO: Hidden irreps: 32x0e 19 | 2024-07-21 14:16:40.986 INFO: Using stochastic weight averaging (after 36 epochs) with energy weight : 1000.0, forces weight : 100.0 and learning rate : 0.001 20 | 2024-07-21 14:16:41.076 INFO: ScaleShiftMACE( 21 | (node_embedding): LinearNodeEmbeddingBlock( 22 | (linear): Linear(3x0e -> 32x0e | 96 weights) 23 | ) 24 | (radial_embedding): RadialEmbeddingBlock( 25 | (bessel_fn): BesselBasis(r_max=4.0, num_basis=8, trainable=False) 26 | (cutoff_fn): PolynomialCutoff(p=5.0, r_max=4.0) 27 | ) 28 | (spherical_harmonics): SphericalHarmonics() 29 | (atomic_energies_fn): AtomicEnergiesBlock(energies=[-10.7072, -48.8474, -102.5712]) 30 | (interactions): ModuleList( 31 | (0): RealAgnosticInteractionBlock( 32 | (linear_up): Linear(32x0e -> 32x0e | 1024 weights) 33 | (conv_tp): TensorProduct(32x0e x 1x0e+1x1o+1x2e -> 32x0e+32x1o+32x2e | 96 paths | 96 weights) 34 | (conv_tp_weights): FullyConnectedNet[8, 64, 64, 64, 96] 35 | (linear): Linear(32x0e+32x1o+32x2e -> 32x0e+32x1o+32x2e | 3072 weights) 36 | (skip_tp): FullyConnectedTensorProduct(32x0e+32x1o+32x2e x 3x0e -> 32x0e+32x1o+32x2e | 9216 paths | 9216 weights) 37 | (reshape): reshape_irreps() 38 | ) 39 | (1): RealAgnosticResidualInteractionBlock( 40 | (linear_up): Linear(32x0e -> 32x0e | 1024 weights) 41 | (conv_tp): TensorProduct(32x0e x 1x0e+1x1o+1x2e -> 32x0e+32x1o+32x2e | 96 paths | 96 weights) 42 | (conv_tp_weights): FullyConnectedNet[8, 64, 64, 64, 96] 43 | (linear): Linear(32x0e+32x1o+32x2e -> 32x0e+32x1o+32x2e | 3072 weights) 44 | (skip_tp): FullyConnectedTensorProduct(32x0e x 3x0e -> 32x0e | 3072 paths | 3072 weights) 45 | (reshape): reshape_irreps() 46 | ) 47 | ) 48 | (products): ModuleList( 49 | (0-1): 2 x EquivariantProductBasisBlock( 50 | (symmetric_contractions): SymmetricContraction( 51 | (contractions): ModuleList( 52 | (0): Contraction( 53 | (contractions_weighting): ModuleList( 54 | (0): GraphModule() 55 | ) 56 | (contractions_features): ModuleList( 57 | (0): GraphModule() 58 | ) 59 | (weights): ParameterList( (0): Parameter containing: [torch.float64 of size 3x1x32 (cuda:0)]) 60 | (graph_opt_main): GraphModule() 61 | ) 62 | ) 63 | ) 64 | (linear): Linear(32x0e -> 32x0e | 1024 weights) 65 | ) 66 | ) 67 | (readouts): ModuleList( 68 | (0): LinearReadoutBlock( 69 | (linear): Linear(32x0e -> 1x0e | 32 weights) 70 | ) 71 | (1): NonLinearReadoutBlock( 72 | (linear_1): Linear(32x0e -> 16x0e | 512 weights) 73 | (non_linearity): Activation [x] (16x0e -> 16x0e) 74 | (linear_2): Linear(16x0e -> 1x0e | 16 weights) 75 | ) 76 | ) 77 | (scale_shift): ScaleShiftBlock(scale=2.177545, shift=0.000000) 78 | ) 79 | 2024-07-21 14:16:41.078 INFO: Number of parameters: 53648 80 | 2024-07-21 14:16:41.078 INFO: Optimizer: Adam ( 81 | Parameter Group 0 82 | amsgrad: True 83 | betas: (0.9, 0.999) 84 | capturable: False 85 | differentiable: False 86 | eps: 1e-08 87 | foreach: None 88 | fused: None 89 | initial_lr: 0.01 90 | lr: 0.01 91 | maximize: False 92 | name: embedding 93 | swa_lr: 0.001 94 | weight_decay: 0.0 95 | 96 | Parameter Group 1 97 | amsgrad: True 98 | betas: (0.9, 0.999) 99 | capturable: False 100 | differentiable: False 101 | eps: 1e-08 102 | foreach: None 103 | fused: None 104 | initial_lr: 0.01 105 | lr: 0.01 106 | maximize: False 107 | name: interactions_decay 108 | swa_lr: 0.001 109 | weight_decay: 5e-07 110 | 111 | Parameter Group 2 112 | amsgrad: True 113 | betas: (0.9, 0.999) 114 | capturable: False 115 | differentiable: False 116 | eps: 1e-08 117 | foreach: None 118 | fused: None 119 | initial_lr: 0.01 120 | lr: 0.01 121 | maximize: False 122 | name: interactions_no_decay 123 | swa_lr: 0.001 124 | weight_decay: 0.0 125 | 126 | Parameter Group 3 127 | amsgrad: True 128 | betas: (0.9, 0.999) 129 | capturable: False 130 | differentiable: False 131 | eps: 1e-08 132 | foreach: None 133 | fused: None 134 | initial_lr: 0.01 135 | lr: 0.01 136 | maximize: False 137 | name: products 138 | swa_lr: 0.001 139 | weight_decay: 5e-07 140 | 141 | Parameter Group 4 142 | amsgrad: True 143 | betas: (0.9, 0.999) 144 | capturable: False 145 | differentiable: False 146 | eps: 1e-08 147 | foreach: None 148 | fused: None 149 | initial_lr: 0.01 150 | lr: 0.01 151 | maximize: False 152 | name: readouts 153 | swa_lr: 0.001 154 | weight_decay: 0.0 155 | ) 156 | 2024-07-21 14:16:41.078 INFO: Using gradient clipping with tolerance=10.000 157 | 2024-07-21 14:16:41.078 INFO: Started training 158 | 2024-07-21 14:16:41.934 INFO: Epoch None: loss=71.3685, RMSE_E_per_atom=6269.3 meV, RMSE_F=2604.0 meV / A 159 | 2024-07-21 14:16:49.644 INFO: Epoch 0: loss=30.7212, RMSE_E_per_atom=5323.5 meV, RMSE_F=1678.3 meV / A 160 | 2024-07-21 14:16:51.989 INFO: Epoch 2: loss=8.4117, RMSE_E_per_atom=4865.6 meV, RMSE_F=779.0 meV / A 161 | 2024-07-21 14:16:53.776 INFO: Epoch 4: loss=4.8551, RMSE_E_per_atom=4090.0 meV, RMSE_F=564.8 meV / A 162 | 2024-07-21 14:16:55.604 INFO: Epoch 6: loss=2.6657, RMSE_E_per_atom=3270.8 meV, RMSE_F=400.3 meV / A 163 | 2024-07-21 14:16:57.407 INFO: Epoch 8: loss=1.8553, RMSE_E_per_atom=1755.6 meV, RMSE_F=394.8 meV / A 164 | 2024-07-21 14:16:59.264 INFO: Epoch 10: loss=1.6977, RMSE_E_per_atom=1053.1 meV, RMSE_F=399.6 meV / A 165 | 2024-07-21 14:17:01.184 INFO: Epoch 12: loss=2.0717, RMSE_E_per_atom=750.7 meV, RMSE_F=449.8 meV / A 166 | 2024-07-21 14:17:03.068 INFO: Epoch 14: loss=1.3686, RMSE_E_per_atom=539.3 meV, RMSE_F=366.5 meV / A 167 | 2024-07-21 14:17:05.032 INFO: Epoch 16: loss=0.8913, RMSE_E_per_atom=354.8 meV, RMSE_F=296.9 meV / A 168 | 2024-07-21 14:17:06.891 INFO: Epoch 18: loss=1.4061, RMSE_E_per_atom=422.7 meV, RMSE_F=372.9 meV / A 169 | 2024-07-21 14:17:08.812 INFO: Epoch 20: loss=1.2417, RMSE_E_per_atom=373.6 meV, RMSE_F=351.1 meV / A 170 | 2024-07-21 14:17:10.867 INFO: Epoch 22: loss=0.7003, RMSE_E_per_atom=267.4 meV, RMSE_F=263.9 meV / A 171 | 2024-07-21 14:17:12.780 INFO: Epoch 24: loss=1.0258, RMSE_E_per_atom=293.0 meV, RMSE_F=319.3 meV / A 172 | 2024-07-21 14:17:14.761 INFO: Epoch 26: loss=0.7541, RMSE_E_per_atom=168.0 meV, RMSE_F=274.7 meV / A 173 | 2024-07-21 14:17:16.613 INFO: Epoch 28: loss=0.8692, RMSE_E_per_atom=177.4 meV, RMSE_F=294.5 meV / A 174 | 2024-07-21 14:17:18.581 INFO: Epoch 30: loss=1.0723, RMSE_E_per_atom=263.1 meV, RMSE_F=326.2 meV / A 175 | 2024-07-21 14:17:20.633 INFO: Epoch 32: loss=0.7172, RMSE_E_per_atom=286.4 meV, RMSE_F=266.1 meV / A 176 | 2024-07-21 14:17:22.545 INFO: Epoch 34: loss=1.2195, RMSE_E_per_atom=177.7 meV, RMSE_F=349.2 meV / A 177 | 2024-07-21 14:17:23.412 INFO: Changing loss based on SWA 178 | 2024-07-21 14:17:24.351 INFO: Epoch 36: loss=0.7185, RMSE_E_per_atom=48.5 meV, RMSE_F=220.0 meV / A 179 | 2024-07-21 14:17:26.285 INFO: Epoch 38: loss=0.5541, RMSE_E_per_atom=29.4 meV, RMSE_F=216.5 meV / A 180 | 2024-07-21 14:17:28.296 INFO: Epoch 40: loss=0.5076, RMSE_E_per_atom=22.4 meV, RMSE_F=213.9 meV / A 181 | 2024-07-21 14:17:30.381 INFO: Epoch 42: loss=0.5004, RMSE_E_per_atom=21.8 meV, RMSE_F=213.0 meV / A 182 | 2024-07-21 14:17:32.258 INFO: Epoch 44: loss=0.4875, RMSE_E_per_atom=9.6 meV, RMSE_F=218.9 meV / A 183 | 2024-07-21 14:17:34.125 INFO: Epoch 46: loss=0.4632, RMSE_E_per_atom=19.0 meV, RMSE_F=206.9 meV / A 184 | 2024-07-21 14:17:36.028 INFO: Epoch 48: loss=0.4545, RMSE_E_per_atom=16.7 meV, RMSE_F=206.8 meV / A 185 | 2024-07-21 14:17:36.951 INFO: Training complete 186 | 2024-07-21 14:17:36.952 INFO: Computing metrics for training, validation, and test sets 187 | 2024-07-21 14:17:37.803 INFO: Loading checkpoint: MACE_models/mace01_run-123_epoch-22.pt 188 | 2024-07-21 14:17:37.828 INFO: Loaded model from epoch 22 189 | 2024-07-21 14:17:37.828 INFO: Evaluating train ... 190 | 2024-07-21 14:17:38.297 INFO: Evaluating valid ... 191 | 2024-07-21 14:17:38.346 INFO: Evaluating Default ... 192 | 2024-07-21 14:17:41.265 INFO: 193 | +-------------+---------------------+------------------+-------------------+ 194 | | config_type | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % | 195 | +-------------+---------------------+------------------+-------------------+ 196 | | train | 294.9 | 245.8 | 11.29 | 197 | | valid | 267.4 | 263.9 | 10.20 | 198 | | Default | 291.2 | 268.8 | 11.74 | 199 | +-------------+---------------------+------------------+-------------------+ 200 | 2024-07-21 14:17:41.265 INFO: Saving model to MACE_models/mace01_run-123.model 201 | 2024-07-21 14:17:41.333 INFO: Compiling model, saving metadata to MACE_models/mace01_compiled.model 202 | 2024-07-21 14:17:41.971 INFO: Loading checkpoint: MACE_models/mace01_run-123_epoch-48_swa.pt 203 | 2024-07-21 14:17:41.998 INFO: Loaded model from epoch 48 204 | 2024-07-21 14:17:41.999 INFO: Evaluating train ... 205 | 2024-07-21 14:17:42.487 INFO: Evaluating valid ... 206 | 2024-07-21 14:17:42.539 INFO: Evaluating Default ... 207 | 2024-07-21 14:17:45.540 INFO: 208 | +-------------+---------------------+------------------+-------------------+ 209 | | config_type | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % | 210 | +-------------+---------------------+------------------+-------------------+ 211 | | train | 19.1 | 182.0 | 8.36 | 212 | | valid | 16.7 | 206.8 | 8.00 | 213 | | Default | 20.4 | 211.4 | 9.23 | 214 | +-------------+---------------------+------------------+-------------------+ 215 | 2024-07-21 14:17:45.541 INFO: Saving model to MACE_models/mace01_run-123_swa.model 216 | 2024-07-21 14:17:45.610 INFO: Compiling model, saving metadata MACE_models/mace01_swa_compiled.model 217 | 2024-07-21 14:17:46.275 INFO: Done 218 | -------------------------------------------------------------------------------- /MACE_models/mace01_swa.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace01_swa.model -------------------------------------------------------------------------------- /MACE_models/mace01_swa_compiled.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace01_swa_compiled.model -------------------------------------------------------------------------------- /MACE_models/mace02_com1.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com1.model -------------------------------------------------------------------------------- /MACE_models/mace02_com1_compiled.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com1_compiled.model -------------------------------------------------------------------------------- /MACE_models/mace02_com1_gen1.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com1_gen1.model -------------------------------------------------------------------------------- /MACE_models/mace02_com1_gen1_compiled.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com1_gen1_compiled.model -------------------------------------------------------------------------------- /MACE_models/mace02_com1_gen1_swa.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com1_gen1_swa.model -------------------------------------------------------------------------------- /MACE_models/mace02_com1_gen1_swa_compiled.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com1_gen1_swa_compiled.model -------------------------------------------------------------------------------- /MACE_models/mace02_com1_swa.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com1_swa.model -------------------------------------------------------------------------------- /MACE_models/mace02_com1_swa_compiled.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com1_swa_compiled.model -------------------------------------------------------------------------------- /MACE_models/mace02_com2.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com2.model -------------------------------------------------------------------------------- /MACE_models/mace02_com2_compiled.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com2_compiled.model -------------------------------------------------------------------------------- /MACE_models/mace02_com2_swa.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com2_swa.model -------------------------------------------------------------------------------- /MACE_models/mace02_com2_swa_compiled.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com2_swa_compiled.model -------------------------------------------------------------------------------- /MACE_models/mace02_com3.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com3.model -------------------------------------------------------------------------------- /MACE_models/mace02_com3_compiled.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com3_compiled.model -------------------------------------------------------------------------------- /MACE_models/mace02_com3_swa.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com3_swa.model -------------------------------------------------------------------------------- /MACE_models/mace02_com3_swa_compiled.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/MACE_models/mace02_com3_swa_compiled.model -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | - [T01-MACE-Practice-I](T01-MACE-Practice-I.ipynb) [![badge](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/imagdau/Tutorials/blob/main/T01-MACE-Practice-I.ipynb) 3 | - [T02-MACE-Practice-II](T02-MACE-Practice-II.ipynb) [![badge](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/imagdau/Tutorials/blob/main/T02-MACE-Practice-II.ipynb) 4 | 5 | - [T03-MACE-Theory](T03-MACE-Theory.ipynb) [![badge](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/imagdau/Tutorials/blob/main/T03-MACE-Theory.ipynb) 6 | 7 | - [T04-MLIP-Apps.ipynb](T04-MLIP-Apps.ipynb) [![badge](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/imagdau/Tutorials/blob/main/T04-MLIP-Apps.ipynb) 8 | 9 | -------------------------------------------------------------------------------- /T01-MACE-Practice-I.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "# MACE in Practice I" 8 | ] 9 | }, 10 | { 11 | "cell_type": "code", 12 | "execution_count": null, 13 | "metadata": {}, 14 | "outputs": [], 15 | "source": [ 16 | "!pip install xtb mace-torch nglview ipywidgets rdkit\n", 17 | "!pip install git+https://github.com/imagdau/aseMolec@main" 18 | ] 19 | }, 20 | { 21 | "cell_type": "markdown", 22 | "metadata": { 23 | "cell_id": "95057ab91ce54145b569a1b35cc81742", 24 | "deepnote_cell_type": "markdown", 25 | "tags": [] 26 | }, 27 | "source": [ 28 | "In this tutorial, you will learn how to fit and test a `MACE` model (Message Passing Neural Network), which is a highly accurate and efficient MLIP (Machine Learnt Interatomic Potential). The training/testing techniques we show here, however, are broadly applicable to all MLIPs. You can independently learn about MACE by studying the [original method paper](https://proceedings.neurips.cc/paper_files/paper/2022/file/4a36c3c51af11ed9f34615b81edb5bbc-Paper-Conference.pdf). MACE was developed by unifying the Atomic Cluster Expansion (ACE) approach with the Neural Equivariant Interatomic Potentials (NequIP). The mathematical formalism which unifies these methods is explained in the [accompaning paper](https://doi.org/10.48550/arXiv.2205.06643). Another [useful reference](https://doi.org/10.48550/arXiv.2305.14247) showcases the method's performance on published benchmark datasets. The [code implementation](https://github.com/ACEsuit/mace) is publically available and [here](https://mace-docs.readthedocs.io/en/latest/) you can find the documentation.\n", 29 | "\n", 30 | "## Learning Objectives for today:\n", 31 | "\n", 32 | "1. **Understanding the data: diverse configs, reference labels**\n", 33 | "2. **Understanding MACE parameters: architecture and training**\n", 34 | "3. **Fitting and testing MACE models**\n", 35 | "4. **Ultimate goal: stable and accurate Molecular Dynamics**" 36 | ] 37 | }, 38 | { 39 | "cell_type": "markdown", 40 | "metadata": {}, 41 | "source": [ 42 | "## 1. Understanding the data" 43 | ] 44 | }, 45 | { 46 | "cell_type": "markdown", 47 | "metadata": { 48 | "jp-MarkdownHeadingCollapsed": true 49 | }, 50 | "source": [ 51 | "### 1.1 Diverse Molecular Conformations" 52 | ] 53 | }, 54 | { 55 | "cell_type": "markdown", 56 | "metadata": {}, 57 | "source": [ 58 | "Understanding the data is a crucial part of fitting an MLIP. Most models will underperform at first, often because of insufficiently representative data.\n", 59 | "In this application, we will develop an MLIP for molecular liquids of carbonates. The data presented here is a subset from [this work](https://doi.org/10.1021/acs.jpcb.2c03746) and comprises a mixture of 6 different types of molecules: cyclic carbonates (Vinylene carbonate VC, Ethylene carbonate EC, Propylene carbonate PC) and linear carbonates (Dimethyl carbonate DMC, Ethyl Methyl Carbonate EMC, Diethyl carbonate DEC). Mixtures of these molecules in various formulations are used as solvents in Li-ion battery electrolytes." 60 | ] 61 | }, 62 | { 63 | "cell_type": "code", 64 | "execution_count": null, 65 | "metadata": {}, 66 | "outputs": [], 67 | "source": [ 68 | "from rdkit import Chem\n", 69 | "from rdkit.Chem import Draw\n", 70 | "\n", 71 | "# SMILES strings for each molecule\n", 72 | "sm_dict = {\n", 73 | " 'VC': 'c1coc(=O)o1',\n", 74 | " 'EC': 'C1COC(=O)O1',\n", 75 | " 'PC': 'CC1COC(=O)O1',\n", 76 | " 'DMC': 'COC(=O)OC',\n", 77 | " 'EMC': 'CCOC(=O)OC',\n", 78 | " 'DEC': 'CCOC(=O)OCC'\n", 79 | "}\n", 80 | "\n", 81 | "Draw.MolsToGridImage([Chem.MolFromSmiles(sm_dict[mol]) for mol in sm_dict], legends=list(sm_dict.keys()))" 82 | ] 83 | }, 84 | { 85 | "cell_type": "markdown", 86 | "metadata": {}, 87 | "source": [ 88 | "For this tutorial, we prepared in advance a collection of atomic configurations (small subset from [this paper](https://doi.org/10.1021/acs.jpcb.2c03746)). Let's understand the data! We start by loading the raw configurations with no `labels` (energy, forces). The atomic `configurations`are stored in the [extxyz](https://wiki.fysik.dtu.dk/ase/ase/io/formatoptions.html#extxyz) format and can be accessed using [ASE](https://wiki.fysik.dtu.dk/ase/index.html) as shown below:" 89 | ] 90 | }, 91 | { 92 | "cell_type": "code", 93 | "execution_count": null, 94 | "metadata": {}, 95 | "outputs": [], 96 | "source": [ 97 | "from ase.io import read, write\n", 98 | "import numpy as np\n", 99 | "\n", 100 | "db = read('data/solvent_configs.xyz', ':') #read in list of configs\n", 101 | "\n", 102 | "print(\"Number of configs in database: \", len(db))\n", 103 | "print(\"Number of atoms in each config: \", np.array([len(at) for at in db]))\n", 104 | "print(\"Number of atoms in the smallest config: \", np.min([len(at) for at in db])) #test if database contains isolated atoms\n", 105 | "print(\"Information stored in config.info: \\n\", db[10].info) #check info\n", 106 | "print(\"Information stored in config.arrays: \\n\", db[10].arrays)" 107 | ] 108 | }, 109 | { 110 | "cell_type": "markdown", 111 | "metadata": {}, 112 | "source": [ 113 | "You can visualize `*.xyz` files using `ovito`. You can launch the `Desktop` application from the `Launcher` and navigate to your working directory in the terminal, then simply run `ovito .xyz`.\n", 114 | "\n", 115 | "At this point, each configuration is a collection of atoms: atomic number (Z) and positions (R), with no additional information. Let's identify the molecules and label molecular clusters. This will make it easier to inspect the data set and, later, test the accuracy of the potential on describing inter-molecular interactions. Molecule identification is achieved using the `wrap_molecs` function from the [aseMolec package](https://github.com/imagdau/aseMolec), shown here for the first 100 frames `db[:100]`." 116 | ] 117 | }, 118 | { 119 | "cell_type": "code", 120 | "execution_count": null, 121 | "metadata": {}, 122 | "outputs": [], 123 | "source": [ 124 | "from aseMolec import anaAtoms as aa\n", 125 | "\n", 126 | "aa.wrap_molecs(db[:100], prog=False) #identify molecules and label molecular clusters, showcase: first 100 frames\n", 127 | "# write('data/solvent_molecs.xyz', db) #save full result\n", 128 | "print(\"Information stored in config.info: \\n\", db[10].info)\n", 129 | "print(\"Information stored in config.arrays: \\n\", db[10].arrays)" 130 | ] 131 | }, 132 | { 133 | "cell_type": "markdown", 134 | "metadata": {}, 135 | "source": [ 136 | "Note the additional information for each atomic config: number of molecules `Nmols`, molecular composition `Comp` (e.g `DEC(1):EC(1)` means the config comprises a dimer with 1 DEC molecule and 1 EC molecule) and molecular ID `molID`. Running the code for the full 5000 configurations can be slow, let's just load the final result (`data/solvent_molecs.xyz`) and inspect the distribution of configs by number of molecules:" 137 | ] 138 | }, 139 | { 140 | "cell_type": "code", 141 | "execution_count": null, 142 | "metadata": {}, 143 | "outputs": [], 144 | "source": [ 145 | "from matplotlib import pyplot as plt\n", 146 | "\n", 147 | "db = read('data/solvent_molecs.xyz', ':')\n", 148 | "Nmols = np.array([at.info['Nmols'] for at in db]) #collect Nmols information across all data\n", 149 | "plt.hist(Nmols, align='left', bins=[1,2,3,4,5,6,7], rwidth=0.8);\n", 150 | "plt.xlabel('# Molecs');\n", 151 | "plt.ylabel('# Configs comprising that # Molecs');" 152 | ] 153 | }, 154 | { 155 | "cell_type": "markdown", 156 | "metadata": {}, 157 | "source": [ 158 | "There are just under 1000 configs comprising of single molecules and more than 2000 dimers. The largest configs contain clusters of six molecules.\n", 159 | "\n", 160 | "We can check the distribution of molecular compositions for each cluster size:" 161 | ] 162 | }, 163 | { 164 | "cell_type": "code", 165 | "execution_count": null, 166 | "metadata": {}, 167 | "outputs": [], 168 | "source": [ 169 | "from aseMolec import extAtoms as ea\n", 170 | "from collections import Counter\n", 171 | "\n", 172 | "comp_dict = {} #create a dictionary of compositions for each cluster size\n", 173 | "for Nmol in range(1,7):\n", 174 | " comp_dict[Nmol] = dict(Counter([at.info['Comp'] for at in ea.sel_by_info_val(db, 'Nmols', Nmol)]))\n", 175 | "\n", 176 | "Nmol = 6 #show distribution of compositions for cluster size 6\n", 177 | "plt.pie(comp_dict[Nmol].values(),\n", 178 | " labels=comp_dict[Nmol].keys(),\n", 179 | " explode=10/(25+np.array(list(comp_dict[Nmol].values()))),\n", 180 | " rotatelabels =True);" 181 | ] 182 | }, 183 | { 184 | "cell_type": "markdown", 185 | "metadata": {}, 186 | "source": [ 187 | "The training set is quite diverse and it contains a good mix of compositions. Check the distribution for other cluster sizes: `Nmol = 1, 2, 3, 4, 5`. Find out if all isolated molecules are present and well sampled. We have six molecules, so there should be 6x7/2 dimers present, are all dimers sampled?" 188 | ] 189 | }, 190 | { 191 | "cell_type": "markdown", 192 | "metadata": { 193 | "jp-MarkdownHeadingCollapsed": true 194 | }, 195 | "source": [ 196 | "### 1.2 Labeling Data with XTB Values" 197 | ] 198 | }, 199 | { 200 | "cell_type": "markdown", 201 | "metadata": {}, 202 | "source": [ 203 | "We convinced ourselves the training set is quite diverse, it samples many compositions and molecular cluster sizes. It is time to prepare the reference data (energies, forces) to train the model on. We will do this using the Semiempirical Tight Binding level of theory with [XTB](https://xtb-docs.readthedocs.io/en/latest/contents.html). This may be less accurate than other methods specialized for these systems, but it is fast and it will later allow us to test MLIP errors on-the-fly.\n", 204 | "\n", 205 | "Notice the data set contains isolated molecules but no isolated atoms. MACE (and other MLIPs) fit to atomization energies (eV) which is total energy minus the energy of each atom in vacuum $(E^{0})$:\n", 206 | "$$\n", 207 | "E^{\\rm atm} = E^{\\rm tot}-\\sum_i^{N} E^{0}\n", 208 | "$$\n", 209 | "\n", 210 | "In our specific example, all molecules comprise of three chemical elements and we will need to compute $(E^{0})$ for each of them:\n", 211 | "\n", 212 | "$$\n", 213 | "E^{\\rm atm} = E^{\\rm tot}-\\sum_i^{N_H} E^{H}_i-\\sum_i^{N_C} E^{C}_i-\\sum_i^{N_O} E^{O}_i\n", 214 | "$$\n", 215 | "\n", 216 | "Let us add three frames containing Hydrogen H, Carbon C and Oxygen O to the dataset and label them as `config_type=IsolatedAtom`" 217 | ] 218 | }, 219 | { 220 | "cell_type": "code", 221 | "execution_count": null, 222 | "metadata": {}, 223 | "outputs": [], 224 | "source": [ 225 | "from ase import Atoms\n", 226 | "\n", 227 | "db = read('data/solvent_molecs.xyz', ':')\n", 228 | "db = [Atoms('H'), Atoms('C'), Atoms('O')]+db #add isolated atoms to the database\n", 229 | "\n", 230 | "for at in db[:3]:\n", 231 | " at.info['config_type'] = 'IsolatedAtom'\n", 232 | "\n", 233 | "print(\"Number of configs in database: \", len(db))" 234 | ] 235 | }, 236 | { 237 | "cell_type": "markdown", 238 | "metadata": {}, 239 | "source": [ 240 | "We are now ready to compute the energy and forces with XTB:" 241 | ] 242 | }, 243 | { 244 | "cell_type": "code", 245 | "execution_count": null, 246 | "metadata": {}, 247 | "outputs": [], 248 | "source": [ 249 | "from tqdm import tqdm\n", 250 | "from xtb.ase.calculator import XTB\n", 251 | "xtb_calc = XTB(method=\"GFN2-xTB\")\n", 252 | "\n", 253 | "for at in tqdm(db[:15]): #showcase: first 15 frames\n", 254 | " at.calc = xtb_calc\n", 255 | " at.info['energy_xtb'] = at.get_potential_energy()\n", 256 | " at.arrays['forces_xtb'] = at.get_forces()\n", 257 | "# write('data/solvent_xtb.xyz', db) #save full result\n", 258 | "\n", 259 | "print(\"Information stored in config.info: \\n\", db[13].info) #check info\n", 260 | "print(\"Information stored in config.arrays: \\n\", db[13].arrays)" 261 | ] 262 | }, 263 | { 264 | "cell_type": "markdown", 265 | "metadata": {}, 266 | "source": [ 267 | "The updated data contains one energy value for each config `energy_xtb` and the `forces_xtb` on each atom. Latest version of [ASE](https://wiki.fysik.dtu.dk/ase/index.html) does not support simple names such as `energy` and `forces` so we append `_xtb`. The entire computation takes about `25 mins` for the 5003 configs. We have precomputed the data, so we can simply load the final result. Let's check the $E^0$ values and atomization energies:" 268 | ] 269 | }, 270 | { 271 | "cell_type": "code", 272 | "execution_count": null, 273 | "metadata": {}, 274 | "outputs": [], 275 | "source": [ 276 | "db = read('data/solvent_xtb.xyz', ':15')\n", 277 | "\n", 278 | "print(\"E0s: \\n\", ea.get_E0(db, tag='_xtb'))\n", 279 | "print(\"Total energy per config: \\n\", ea.get_prop(db, 'info', 'energy_xtb', peratom=False)[13])\n", 280 | "print(\"Toal energy per atom: \\n\", ea.get_prop(db, 'info', 'energy_xtb', peratom=True)[13])\n", 281 | "print(\"Atomization energy per config: \\n\", ea.get_prop(db, 'bind', prop='_xtb', peratom=False)[13])\n", 282 | "print(\"Atomization energy per atom: \\n\", ea.get_prop(db, 'bind', prop='_xtb', peratom=True)[13])" 283 | ] 284 | }, 285 | { 286 | "cell_type": "markdown", 287 | "metadata": {}, 288 | "source": [ 289 | "Good! We get about -6 $\\rm eV/atom$ which is largely dominated by the energy of the [covalent bonds](https://en.wikipedia.org/wiki/Bond-dissociation_energy#:~:text=of%20a%20solvent.-,Representative%20bond%20enthalpies,-%5Bedit%5D). Remember, the largest contribution to the total energy comes from $E^0$ and then from covalent bonds. The noncovalent interactions contribute a very small amount to the total energy, yet they are crucial for molecular dynamics." 290 | ] 291 | }, 292 | { 293 | "cell_type": "markdown", 294 | "metadata": { 295 | "cell_id": "6e1b8f3b26c64c69982a15b403550fdb", 296 | "deepnote_cell_type": "markdown" 297 | }, 298 | "source": [ 299 | "## 2. Understanding MACE parameters" 300 | ] 301 | }, 302 | { 303 | "cell_type": "markdown", 304 | "metadata": { 305 | "jp-MarkdownHeadingCollapsed": true 306 | }, 307 | "source": [ 308 | "### 2.1 Model parameters" 309 | ] 310 | }, 311 | { 312 | "cell_type": "markdown", 313 | "metadata": {}, 314 | "source": [ 315 | "We'll give a high-level explanation of the important parameters in MACE. We will discuss these in detail during the third tutorial and associated lectures. Consult the [documentation](https://github.com/ACEsuit/mace) for additional parameters.\n", 316 | "\n", 317 | "- ##### --num_interactions: message-passing layers\n", 318 | "\n", 319 | "Controls the number of message-passing layers in the model.\n", 320 | "\n", 321 | "- ##### --hidden_irreps: number of message passing layers\n", 322 | "\n", 323 | "Determines the size of the model and its symmetry.\n", 324 | "For example: `hidden_irreps='128x0e'` means the model has `128` channels or paths, the output is invariant under rotation ($L_{\\rm max}=0$) and even under inversion (`'e'`). For most applications, these settings will do well. `hidden_irreps='64x0e + 64x1o'` means the model has `64` channels and is equivariant under rotation ($L_{\\rm max}=0$).\n", 325 | "\n", 326 | "Alternatively, the model size can be adjusted using a pair of more user-friendly arguments:\n", 327 | "\n", 328 | "           --num_channels=32\n", 329 | "\n", 330 | "           --max_L=2\n", 331 | "\n", 332 | "which, taken together achieve the same as --hidden_irreps='32x0e + 32x1o + 32x2e'\n", 333 | " \n", 334 | "**In general, the `accuracy` of the model can be improved by using more layers, more channels or higher equivariances. This will result in more parameters and `slower` models.**\n", 335 | "\n", 336 | "- ##### --correlation: the order of the many-body expansion\n", 337 | "$$\n", 338 | "E_{i} = E^{(0)}_{i} + \\sum_{j} E_{ij}^{(1)} + \\sum_{jk} E_{ijk}^{(2)} + ...\n", 339 | "$$\n", 340 | "\n", 341 | "The energy-expansion order that MACE induces at each layer. Choosing `--correlation=3` will create basis functions of up to 4-body (ijkl) indices, for each layer. If the model has multiple layers, the effective correlation order is higher. For example, a two-layer MACE with `--correlation=3` has an effective body order of `13`.\n", 342 | "\n", 343 | "- ##### --r_max: the cutoff radius\n", 344 | "\n", 345 | "The cut-off applied to local environment in each layer. `r_max=3.0` means atoms separated by a distance of more than 3.0 A do not directly `communicate`. When the model has multiple message-passing layers, atoms further than 3.0 A can still `communicate` through later messages if intermediate proxy atoms exist. The effective `receptive field` of the model is `num_interactions x r_max`.\n", 346 | "\n", 347 | "- ##### --max_ell: angular resolution\n", 348 | "\n", 349 | "The angular resolution describes how well the model can describe angles. This is controlled by `l_max` of the spherical harmonics basis (not to be confused with `L_max`). Larger values will result in more accurate but slower models. The default is `l_max=3`, appropriate in most cases.\n", 350 | "\n", 351 | "Let's train our first model:" 352 | ] 353 | }, 354 | { 355 | "cell_type": "code", 356 | "execution_count": null, 357 | "metadata": {}, 358 | "outputs": [], 359 | "source": [ 360 | "%%writefile config/config-01.yml\n", 361 | "\n", 362 | "model: MACE\n", 363 | "num_interactions: 2 \n", 364 | "num_channels: 32\n", 365 | "max_L: 0\n", 366 | "correlations: 2 \n", 367 | "r_max: 4.0\n", 368 | "max_ell: 2" 369 | ] 370 | }, 371 | { 372 | "cell_type": "code", 373 | "execution_count": null, 374 | "metadata": {}, 375 | "outputs": [], 376 | "source": [ 377 | "!mace_run_train --config config/config-01.yaml --seed 42" 378 | ] 379 | }, 380 | { 381 | "cell_type": "markdown", 382 | "metadata": {}, 383 | "source": [ 384 | "This won't work, of course, we still need to supply the training data and a model name, at a minimum. In fact, let's take a look at the remaining parameters which control the file management and training protocol." 385 | ] 386 | }, 387 | { 388 | "cell_type": "markdown", 389 | "metadata": { 390 | "jp-MarkdownHeadingCollapsed": true 391 | }, 392 | "source": [ 393 | "### 2.2 Training and data management parameters" 394 | ] 395 | }, 396 | { 397 | "cell_type": "markdown", 398 | "metadata": {}, 399 | "source": [ 400 | "Let's start by splitting the data into a train and test set." 401 | ] 402 | }, 403 | { 404 | "cell_type": "code", 405 | "execution_count": null, 406 | "metadata": { 407 | "cell_id": "ccb3270bdaab4f2490eb8d7d417a260c", 408 | "deepnote_cell_type": "code", 409 | "deepnote_to_be_reexecuted": false, 410 | "execution_millis": 4511, 411 | "execution_start": 1690385915880, 412 | "source_hash": null 413 | }, 414 | "outputs": [], 415 | "source": [ 416 | "from ase.io import read, write\n", 417 | "\n", 418 | "db = read('data/solvent_xtb.xyz', ':')\n", 419 | "write('data/solvent_xtb_train_200.xyz', db[:203]) #first 200 configs plus the 3 E0s\n", 420 | "write('data/solvent_xtb_test.xyz', db[-1000:]) #last 1000 configs" 421 | ] 422 | }, 423 | { 424 | "cell_type": "markdown", 425 | "metadata": { 426 | "cell_id": "efb4e7559b24449d8ec15b35cf8ad561", 427 | "deepnote_cell_type": "markdown" 428 | }, 429 | "source": [ 430 | "- ##### --name: the name of the model\n", 431 | "This name will be used to form file names (model, log, checkpoints, results), so choose a distinct name for each experiment\n", 432 | "\n", 433 | "- ##### --model_dir, --log_dir, --checkpoints_dir, --results_dir: directory paths\n", 434 | "These are the directories where each type of file is saved. For simplicity, we will all files in the same directory.\n", 435 | "\n", 436 | "- ##### --train_file: name of training data\n", 437 | "\n", 438 | "These data configs are used to compute gradients and update model parameters.\n", 439 | "\n", 440 | "- ##### --valid_file: name of validation data\n", 441 | "An alternative way to choose the validation set is by using the `--valid_fraction` keyword. These data configs are used to estimate the model accuracy during training, but not for parameter optimization. The validation set also controls the stopping of the training. At each `--eval_interval` the model is tested on the validation set. The evaluation of these configs takes place in batches, which can be controlled by `--valid_batch_size`. If the accuracy of the model stops improving on the validation set for `--patience` number of epochs, the model will undergo **early stopping**. \n", 442 | "\n", 443 | "- ##### --test_file: name of testing data\n", 444 | "\n", 445 | "This set is entirely independent and only gets evaluated at the end of the training process to estimate the model accuracy on an independent set.\n", 446 | "\n", 447 | "- ##### --E0s: isolated atom energies\n", 448 | "\n", 449 | "Controls how `E0s` should be determined. The strongly recommended approach is to add these values to the training set with `config_type=IsolatedAtom` in `atoms.info` and set `E0s=\"isolated\"`. If these values are not available, MACE can estimate them by least square regression over the available data `E0s=\"average\"` which can lead to unintended consequences depending on how representative the data is.\n", 450 | "\n", 451 | "- ##### --energy_key, --forces_key the key where these values are stores\n", 452 | "This key must coincide with the `ase.Atoms.info[key]/ase.Atoms.arrays[key]` where the energies and forces are stored in the ase.Atoms object.\n", 453 | "\n", 454 | "- ##### --device computing device to use\n", 455 | "Can be CPU (`cpu`), GPU (`cuda`) or Apple Silicon (`mps`). Here we will use `cuda` since the GPU will be significantly faster than the CPU.\n", 456 | "\n", 457 | "- ##### --batch_size number of configs evaluated in one batch\n", 458 | "Number of configs used to compute the gradients for each full update of the network parameters. This training strategy is called stochastic gradient descent because only a subset of the data (`batch_size`) is used to change the parameters at each update.\n", 459 | "\n", 460 | "- ##### --max_num_epochs number of passes through the data\n", 461 | "An `epoch` is completed when the entire training data has been used once in updating the weights `batch` by `batch`. A new epoch begins, and the process repeats.\n", 462 | "\n", 463 | "- ##### --swa protocol for loss weights\n", 464 | "During training you will notice energy errors are at first much higher than force errors, MACE implements a special protocol that increases the weight on the energy in the loss function (`--swa_energy_weight`) once the forces are sufficiently accurate. The starting epoch for this special protocol can be controlled by changing `--start_swa`.\n", 465 | "\n", 466 | "- ##### --seed random number generator seed\n", 467 | "Useful for preparing committee of models.\n", 468 | "\n", 469 | "Now we are ready to fit our first MACE model:" 470 | ] 471 | }, 472 | { 473 | "cell_type": "markdown", 474 | "metadata": { 475 | "cell_id": "01d2525548a64079b8c6247362360f3a", 476 | "deepnote_cell_type": "markdown" 477 | }, 478 | "source": [ 479 | "## 3. Fitting and Testing MACE models" 480 | ] 481 | }, 482 | { 483 | "cell_type": "markdown", 484 | "metadata": { 485 | "jp-MarkdownHeadingCollapsed": true 486 | }, 487 | "source": [ 488 | "### 3.1 Fitting the model" 489 | ] 490 | }, 491 | { 492 | "cell_type": "code", 493 | "execution_count": null, 494 | "metadata": {}, 495 | "outputs": [], 496 | "source": [ 497 | "%%writefile config/config-02.yml\n", 498 | " \n", 499 | "model: \"MACE\" \n", 500 | "num_interactions: 2 \n", 501 | "num_channels: 32 \n", 502 | "max_L: 0 \n", 503 | "correlation: 2 \n", 504 | "r_max: 4.0 \n", 505 | "max_ell: 2 \n", 506 | "name: \"mace01\" \n", 507 | "model_dir: \"MACE_models\" \n", 508 | "log_dir: \"MACE_models\" \n", 509 | "checkpoints_dir: \"MACE_models\" \n", 510 | "results_dir: \"MACE_models\" \n", 511 | "train_file: \"data/solvent_xtb_train_200.xyz\" \n", 512 | "valid_fraction: 0.10 \n", 513 | "test_file: \"data/solvent_xtb_test.xyz\" \n", 514 | "E0s: \"average\" \n", 515 | "energy_key: \"energy_xtb\" \n", 516 | "forces_key: \"forces_xtb\" \n", 517 | "device: cuda \n", 518 | "batch_size: 10 \n", 519 | "max_num_epochs: 50 \n", 520 | "swa: True \n", 521 | "seed: 123\n" 522 | ] 523 | }, 524 | { 525 | "cell_type": "code", 526 | "execution_count": null, 527 | "metadata": {}, 528 | "outputs": [], 529 | "source": [ 530 | "!mace_run_train --config config/config-02.yml\n", 531 | "\n", 532 | "#remove checkpoints since they may cause errors on retraining a model with the same name but a different architecture\n", 533 | "import glob\n", 534 | "import os\n", 535 | "for file in glob.glob(\"MACE_models/*_run-*.model\"):\n", 536 | " os.remove(file)\n", 537 | "for file in glob.glob(\"MACE_models/*.pt\"):\n", 538 | " os.remove(file)" 539 | ] 540 | }, 541 | { 542 | "cell_type": "markdown", 543 | "metadata": { 544 | "cell_id": "b7d69e0612d146dd8c48b408ad4f59a5", 545 | "deepnote_cell_type": "code", 546 | "deepnote_to_be_reexecuted": false, 547 | "execution_millis": 956487, 548 | "execution_start": 1690385920385, 549 | "source_hash": null 550 | }, 551 | "source": [ 552 | "Monitor the memory and usage of your GPU using `nvitop` in the `terminal`. What changes when increasing the `batch_size`? How about when changing the model size (`num_interactions`, `num_channels`, `max_L`, `correlation`)? Notice the accuracy vs speed trade-off.\n", 553 | "\n", 554 | "##### **Question: Here we trained on a very small subset of the data, repeat training for 400, 1000, 2000, 4000 data points, remember to change the name of the model accordingly. How does the learning curve (test error vs size of training test) look like?** " 555 | ] 556 | }, 557 | { 558 | "cell_type": "markdown", 559 | "metadata": { 560 | "cell_id": "9ecbfaac536f4a2bb121e337ca964f99", 561 | "deepnote_cell_type": "markdown" 562 | }, 563 | "source": [ 564 | "We will now use the `mace_eval_configs` script to evaluate the trained model on both the train and test datasets. The script takes the arguments: `--configs` which specifies the file to evaluate, the path to the model in `--model` and the path to the output in `--output`." 565 | ] 566 | }, 567 | { 568 | "cell_type": "markdown", 569 | "metadata": { 570 | "jp-MarkdownHeadingCollapsed": true 571 | }, 572 | "source": [ 573 | "### 3.2 Testing the model: simple RMSEs" 574 | ] 575 | }, 576 | { 577 | "cell_type": "code", 578 | "execution_count": null, 579 | "metadata": { 580 | "cell_id": "f83aa3ea57d241d4adb6d8565648500e", 581 | "deepnote_cell_type": "code", 582 | "deepnote_to_be_reexecuted": false, 583 | "execution_millis": 30282, 584 | "execution_start": 1690386876771, 585 | "source_hash": null 586 | }, 587 | "outputs": [], 588 | "source": [ 589 | "import warnings\n", 590 | "warnings.filterwarnings(\"ignore\")\n", 591 | "\n", 592 | "os.makedirs(\"tests/mace01/\", exist_ok=True)\n", 593 | "\n", 594 | "#evaluate the train set\n", 595 | "!mace_eval_configs \\\n", 596 | " --configs=\"data/solvent_xtb_train_200.xyz\" \\\n", 597 | " --model=\"MACE_models/mace01_swa_compiled.model\" \\\n", 598 | " --output=\"tests/mace01/solvent_train.xyz\"\n", 599 | "\n", 600 | "#evaluate the test set\n", 601 | "!mace_eval_configs \\\n", 602 | " --configs=\"data/solvent_xtb_test.xyz\" \\\n", 603 | " --model=\"MACE_models/mace01_swa_compiled.model\" \\\n", 604 | " --output=\"tests/mace01/solvent_test.xyz\"" 605 | ] 606 | }, 607 | { 608 | "cell_type": "markdown", 609 | "metadata": { 610 | "cell_id": "9863af93365b4659a0fec0989ee10dff", 611 | "deepnote_cell_type": "markdown" 612 | }, 613 | "source": [ 614 | "We can compare MACE vs XTB accuracy on the train and test sets and for this we will use the [aseMolec](git@github.com:imagdau/aseMolec.git) which implements some handy utilities for manipulating ase.Atoms and testing potentials, especially for molecular systems." 615 | ] 616 | }, 617 | { 618 | "cell_type": "code", 619 | "execution_count": null, 620 | "metadata": { 621 | "cell_id": "79b6d4c1b7da48a0b3cfdb0fc1badd17", 622 | "deepnote_cell_type": "code", 623 | "deepnote_to_be_reexecuted": false, 624 | "execution_millis": 4173, 625 | "execution_start": 1690386907048, 626 | "source_hash": null 627 | }, 628 | "outputs": [], 629 | "source": [ 630 | "from aseMolec import pltProps as pp\n", 631 | "from ase.io import read\n", 632 | "import matplotlib.pyplot as plt\n", 633 | "from aseMolec import extAtoms as ea\n", 634 | "import numpy as np\n", 635 | "\n", 636 | "def plot_RMSEs(db, labs):\n", 637 | " ea.rename_prop_tag(db, 'MACE_energy', 'energy_mace') #Backward compatibility\n", 638 | " ea.rename_prop_tag(db, 'MACE_forces', 'forces_mace') #Backward compatibility\n", 639 | " \n", 640 | " plt.figure(figsize=(9,6), dpi=100)\n", 641 | " plt.subplot(1,3,1)\n", 642 | " pp.plot_prop(ea.get_prop(db, 'bind', '_xtb', True).flatten(), \\\n", 643 | " ea.get_prop(db, 'bind', '_mace', True).flatten(), \\\n", 644 | " title=r'Energy $(\\rm eV/atom)$ ', labs=labs, rel=False)\n", 645 | " plt.subplot(1,3,2)\n", 646 | " pp.plot_prop(ea.get_prop(db, 'info', 'energy_xtb', True).flatten(), \\\n", 647 | " ea.get_prop(db, 'info', 'energy_mace', True).flatten(), \\\n", 648 | " title=r'Energy $(\\rm eV/atom)$ ', labs=labs, rel=False)\n", 649 | " plt.subplot(1,3,3)\n", 650 | " pp.plot_prop(np.concatenate(ea.get_prop(db, 'arrays', 'forces_xtb')).flatten(), \\\n", 651 | " np.concatenate(ea.get_prop(db, 'arrays', 'forces_mace')).flatten(), \\\n", 652 | " title=r'Forces $\\rm (eV/\\AA)$ ', labs=labs, rel=False)\n", 653 | " plt.tight_layout()\n", 654 | " return\n", 655 | "\n", 656 | "train_data = read('tests/mace01/solvent_train.xyz', ':')\n", 657 | "test_data = train_data[:3]+read('tests/mace01/solvent_test.xyz', ':') #append the E0s for computing atomization energy errors\n", 658 | "\n", 659 | "plot_RMSEs(train_data, labs=['XTB', 'MACE'])\n", 660 | "plot_RMSEs(test_data, labs=['XTB', 'MACE'])" 661 | ] 662 | }, 663 | { 664 | "cell_type": "markdown", 665 | "metadata": { 666 | "cell_id": "de7f02db063640f1ac3e94c90c8772c2", 667 | "deepnote_cell_type": "markdown" 668 | }, 669 | "source": [ 670 | "These figures show correlation plots between XTB values and MACE predicted values for atomization energy per atom, the total energy per atom and forces. Do the RMSE values match the number printed at the end of the model training? These errors don't look too bad, and this MACE is a small model with few parameters. Significantly better accuracies can be achieved when training on larger models with more data. How does your model trained on 4000 configs compare?" 671 | ] 672 | }, 673 | { 674 | "cell_type": "markdown", 675 | "metadata": { 676 | "jp-MarkdownHeadingCollapsed": true 677 | }, 678 | "source": [ 679 | "### 3.3 Testing on the Intra/Inter decomposition:" 680 | ] 681 | }, 682 | { 683 | "cell_type": "markdown", 684 | "metadata": {}, 685 | "source": [ 686 | "As shown in this [paper](https://doi.org/10.1038/s41524-023-01100-w) one of the challenges associated with modelling molecular systems has to do with the inter-molecular interactions.\n", 687 | "Molecular dynamics is primarily driven by these inter-molecular interactions, however they are relatively small in comparison to covalent interactions and prove difficult to capture with MLIPs.\n", 688 | "The paper introduces this protocol to decompose the force errors into [intra- and inter-] molecular RMSEs to gauge the quality of the potential separately on the two interaction scales. This approach can be summarized as follows:\n", 689 | "1. Identify molecules (labeled **j**). \n", 690 | "\n", 691 | "2. Within each molecule **j** sum over all atomic forces (labeled **k**) to obtain the **translational** component:\n", 692 | "$$F^{\\rm trans}_j = \\sum_{k \\in j} f_{k}$$\n", 693 | "\n", 694 | "3. Redistribute the molecular **translational** force onto individual atoms (labeled **i**) to obtain the atomic **translational** contributions:\n", 695 | "$$f^{\\rm trans}_i = \\frac{m_i}{M_j} F^{\\rm trans}_j$$\n", 696 | "\n", 697 | "4. Similarly, compute the torque on the entire molecule:\n", 698 | "$$T_j = \\sum_{k \\in j} f_{k} \\times r_{k}$$\n", 699 | "\n", 700 | "5. Compute the atomic **rotational** force contributions that give rise to the given molecular torque:\n", 701 | "$$f^{\\rm rot}_i = m_i r_i \\times (I_j^{\\alpha \\beta})^{-1} T_j$$\n", 702 | "\n", 703 | "6. Finally compute the **vibrational** contribution as the difference:\n", 704 | "$$f^{\\rm vib}_i = f_i - f^{\\rm trans}_i - f^{\\rm rot}_i$$\n", 705 | "\n", 706 | "7. In this approach, the **Inter** is the sum of `trans` and `rot`, while the **Intra** is the `vib` component.\n", 707 | "\n", 708 | "This force decomposition can be automatically obtained using the [aseMolec](https://github.com/imagdau/aseMolec) package as shown below:" 709 | ] 710 | }, 711 | { 712 | "cell_type": "code", 713 | "execution_count": null, 714 | "metadata": { 715 | "cell_id": "58506e689b5c487f91bfb60c2be0bd8f", 716 | "deepnote_cell_type": "code", 717 | "deepnote_to_be_reexecuted": false, 718 | "execution_millis": 6188, 719 | "execution_start": 1690386911122, 720 | "source_hash": null 721 | }, 722 | "outputs": [], 723 | "source": [ 724 | "from aseMolec import pltProps as pp\n", 725 | "from aseMolec import anaAtoms as aa\n", 726 | "\n", 727 | "db1 = read('tests/mace01/solvent_test.xyz', ':')\n", 728 | "ea.rename_prop_tag(db1, 'energy_xtb', 'energy') #Backward compatibility\n", 729 | "ea.rename_prop_tag(db1, 'forces_xtb', 'forces') #Backward compatibility\n", 730 | "\n", 731 | "db2 = read('tests/mace01/solvent_test.xyz', ':')\n", 732 | "ea.rename_prop_tag(db2, 'MACE_energy', 'energy') #Backward compatibility\n", 733 | "ea.rename_prop_tag(db2, 'MACE_forces', 'forces') #Backward compatibility\n", 734 | "\n", 735 | "aa.extract_molecs(db1, intra_inter=True)\n", 736 | "aa.extract_molecs(db2, intra_inter=True)\n", 737 | "\n", 738 | "pp.plot_trans_rot_vib(db1, db2, labs=['XTB', 'MACE'])" 739 | ] 740 | }, 741 | { 742 | "cell_type": "markdown", 743 | "metadata": { 744 | "cell_id": "e69c748a1ff34efb9d513bbfd0716c36", 745 | "deepnote_cell_type": "markdown" 746 | }, 747 | "source": [ 748 | "Indeed, the translation and rotational part of the forces (related to inter-molecular interactions) is significantly harder to capture as evidenced by the larger errors. While the absolute RMSEs are smaller, the relative RMSEs are signifincatly larger for the inter-molecular components. Nevertheless, MACE errors ar significantly lower than other models on these tests." 749 | ] 750 | }, 751 | { 752 | "cell_type": "markdown", 753 | "metadata": { 754 | "cell_id": "e7212170b1574fc3995aa2d51560a425", 755 | "deepnote_cell_type": "markdown" 756 | }, 757 | "source": [ 758 | "## 4. Molecular Dynamics with MACE" 759 | ] 760 | }, 761 | { 762 | "cell_type": "markdown", 763 | "metadata": { 764 | "jp-MarkdownHeadingCollapsed": true 765 | }, 766 | "source": [ 767 | "### 4.1 Is the dynamics stable?" 768 | ] 769 | }, 770 | { 771 | "cell_type": "markdown", 772 | "metadata": { 773 | "cell_id": "f8db4a6a266c490681977db3f868e85e", 774 | "deepnote_cell_type": "markdown" 775 | }, 776 | "source": [ 777 | "Accuracy on fixed test sets is great, but molecular dynamics (MD) is the ultimate test. First, we care about stability, then accuracy: let's check if MACE gives stable dynamics. We will start by implementing a simple function to run Langevin dynamics. We initialize the temperature at 300 K and remove all translations and rotations." 778 | ] 779 | }, 780 | { 781 | "cell_type": "code", 782 | "execution_count": null, 783 | "metadata": { 784 | "cell_id": "0ee5f6a5559d43d5b5fa0bf9101b9170", 785 | "deepnote_cell_type": "code", 786 | "deepnote_to_be_reexecuted": false, 787 | "execution_millis": 130, 788 | "execution_start": 1690386918077, 789 | "source_hash": null 790 | }, 791 | "outputs": [], 792 | "source": [ 793 | "from ase.io import read, write\n", 794 | "from ase import units\n", 795 | "from ase.md.langevin import Langevin\n", 796 | "from ase.md.velocitydistribution import Stationary, ZeroRotation, MaxwellBoltzmannDistribution\n", 797 | "\n", 798 | "import random\n", 799 | "import os\n", 800 | "import time\n", 801 | "import numpy as np\n", 802 | "import pylab as pl\n", 803 | "from IPython import display\n", 804 | "\n", 805 | "def simpleMD(init_conf, temp, calc, fname, s, T):\n", 806 | " init_conf.set_calculator(calc)\n", 807 | "\n", 808 | " #initialize the temperature\n", 809 | " random.seed(701) #just making sure the MD failure is reproducible\n", 810 | " MaxwellBoltzmannDistribution(init_conf, temperature_K=300) #initialize temperature at 300\n", 811 | " Stationary(init_conf)\n", 812 | " ZeroRotation(init_conf)\n", 813 | "\n", 814 | " dyn = Langevin(init_conf, 1.0*units.fs, temperature_K=temp, friction=0.1) #drive system to desired temperature\n", 815 | "\n", 816 | " %matplotlib inline\n", 817 | "\n", 818 | " time_fs = []\n", 819 | " temperature = []\n", 820 | " energies = []\n", 821 | "\n", 822 | " #remove previously stored trajectory with the same name\n", 823 | " os.system('rm -rfv '+fname)\n", 824 | "\n", 825 | " fig, ax = pl.subplots(2, 1, figsize=(6,6), sharex='all', gridspec_kw={'hspace': 0, 'wspace': 0})\n", 826 | "\n", 827 | " def write_frame():\n", 828 | " dyn.atoms.write(fname, append=True)\n", 829 | " time_fs.append(dyn.get_time()/units.fs)\n", 830 | " temperature.append(dyn.atoms.get_temperature())\n", 831 | " energies.append(dyn.atoms.get_potential_energy()/len(dyn.atoms))\n", 832 | "\n", 833 | " ax[0].plot(np.array(time_fs), np.array(energies), color=\"b\")\n", 834 | " ax[0].set_ylabel('E (eV/atom)')\n", 835 | "\n", 836 | " # plot the temperature of the system as subplots\n", 837 | " ax[1].plot(np.array(time_fs), temperature, color=\"r\")\n", 838 | " ax[1].set_ylabel('T (K)')\n", 839 | " ax[1].set_xlabel('Time (fs)')\n", 840 | "\n", 841 | " display.clear_output(wait=True)\n", 842 | " display.display(pl.gcf())\n", 843 | " time.sleep(0.01)\n", 844 | "\n", 845 | " dyn.attach(write_frame, interval=s)\n", 846 | " t0 = time.time()\n", 847 | " dyn.run(T)\n", 848 | " t1 = time.time()\n", 849 | " print(\"MD finished in {0:.2f} minutes!\".format((t1-t0)/60))" 850 | ] 851 | }, 852 | { 853 | "cell_type": "markdown", 854 | "metadata": { 855 | "cell_id": "e73034e343594dbfb0991290c0668b3d", 856 | "deepnote_cell_type": "markdown" 857 | }, 858 | "source": [ 859 | "Now we can run MD with MACE and compare it to the XTB dynamics. Let's try 2 picoseoncds at 1200 K, starting from a single molecule config:" 860 | ] 861 | }, 862 | { 863 | "cell_type": "code", 864 | "execution_count": null, 865 | "metadata": { 866 | "cell_id": "d8629178db814314a7aabb2a0bdd2487", 867 | "deepnote_cell_type": "code", 868 | "deepnote_to_be_reexecuted": false, 869 | "execution_millis": 516302, 870 | "execution_start": 1690386918208, 871 | "source_hash": null 872 | }, 873 | "outputs": [], 874 | "source": [ 875 | " #let us start with a single molecule\n", 876 | "init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n", 877 | "\n", 878 | "#we can use MACE as a calculator in ASE!\n", 879 | "from mace.calculators import MACECalculator\n", 880 | "mace_calc = MACECalculator(model_paths=['MACE_models/mace01_swa_compiled.model'], device='cuda', default_dtype=\"float32\")\n", 881 | "\n", 882 | "simpleMD(init_conf, temp=1200, calc=mace_calc, fname='moldyn/mace01_md.xyz', s=10, T=2000)\n" 883 | ] 884 | }, 885 | { 886 | "cell_type": "markdown", 887 | "metadata": {}, 888 | "source": [ 889 | "Let's visualize the trajectory using `nglviewer`. Alternatively you can use `ovito` through the `Launcher`." 890 | ] 891 | }, 892 | { 893 | "cell_type": "code", 894 | "execution_count": null, 895 | "metadata": {}, 896 | "outputs": [], 897 | "source": [ 898 | "import nglview as nv\n", 899 | "traj = read('moldyn/mace01_md.xyz', ':')\n", 900 | "nv.show_asetraj(traj)" 901 | ] 902 | }, 903 | { 904 | "cell_type": "markdown", 905 | "metadata": { 906 | "cell_id": "0ddf1be382f64c4b98e105ab40c77054", 907 | "deepnote_cell_type": "markdown" 908 | }, 909 | "source": [ 910 | "For reference, we can also run XTB dynamics from the same starting configuration." 911 | ] 912 | }, 913 | { 914 | "cell_type": "code", 915 | "execution_count": null, 916 | "metadata": { 917 | "allow_embed": false, 918 | "cell_id": "e3e0c02b2c9942ae88c8c228dbca471f", 919 | "deepnote_cell_type": "code", 920 | "deepnote_to_be_reexecuted": false, 921 | "execution_millis": 1369784, 922 | "execution_start": 1690388047354, 923 | "source_hash": null 924 | }, 925 | "outputs": [], 926 | "source": [ 927 | "# reinitialize the original config\n", 928 | "init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n", 929 | "\n", 930 | "from xtb.ase.calculator import XTB\n", 931 | "xtb_calc = XTB(method=\"GFN2-xTB\")\n", 932 | "\n", 933 | "simpleMD(init_conf, temp=1200, calc=xtb_calc, fname='moldyn/xtb_md.xyz', s=10, T=2000)" 934 | ] 935 | }, 936 | { 937 | "cell_type": "markdown", 938 | "metadata": {}, 939 | "source": [ 940 | "MACE dynamics finished in 2 minutes, vs XTB in 14. This is the essence of MLIP: speeding up calculations that would otherwise take a long time to run. In this case, the speed-up is just one order of magnitude, but depending on the cost of the reference calculation, it can be many orders of magnitude for expensive Quantum Chemistry methods and large systems. Remember, the cost of MLIP is independent of the accuracy of the potential energy surface!\n", 941 | "\n", 942 | "Let's visualize the trajectory:" 943 | ] 944 | }, 945 | { 946 | "cell_type": "code", 947 | "execution_count": null, 948 | "metadata": {}, 949 | "outputs": [], 950 | "source": [ 951 | "import nglview as nv\n", 952 | "traj = read('moldyn/xtb_md.xyz', ':')\n", 953 | "nv.show_asetraj(traj)" 954 | ] 955 | }, 956 | { 957 | "cell_type": "markdown", 958 | "metadata": {}, 959 | "source": [ 960 | "Remember, you can visualize the trajectory files `.xyz` with `ovito` in the `Desktop` opened through the `Launcher`.\n", 961 | "\n", 962 | "Obtaining stable dynamics with so little training is a great result. Up until recently, most MLIPs would require a lot of training before MD was stable. MACE combines the lessons learned over 10-15 years in MLIP development, to achieve a smooth and regular potential energy surface, which minimizes the risk of unstable MD." 963 | ] 964 | }, 965 | { 966 | "cell_type": "markdown", 967 | "metadata": { 968 | "jp-MarkdownHeadingCollapsed": true 969 | }, 970 | "source": [ 971 | "### 4.2 Is the dynamics accurate?" 972 | ] 973 | }, 974 | { 975 | "cell_type": "markdown", 976 | "metadata": { 977 | "cell_id": "b5e16ad359e740ff82138976294bc9dc", 978 | "deepnote_cell_type": "markdown" 979 | }, 980 | "source": [ 981 | "Are the different dynamics sampling the correct distributions? Let us check the radial distribution functions (RDF). The [aseMolec](https://github.com/imagdau/aseMolec) package provides functionality to do that:" 982 | ] 983 | }, 984 | { 985 | "cell_type": "code", 986 | "execution_count": null, 987 | "metadata": { 988 | "cell_id": "860874d8f7b64d08851ff489a42228eb", 989 | "deepnote_cell_type": "code", 990 | "deepnote_to_be_reexecuted": false, 991 | "execution_millis": 23493, 992 | "execution_start": 1690389417019, 993 | "source_hash": null 994 | }, 995 | "outputs": [], 996 | "source": [ 997 | "from aseMolec import anaAtoms as aa\n", 998 | "\n", 999 | "tag = 'HO_intra' #choose one of 'HH_intra', 'HC_intra', 'HO_intra', 'CC_intra', 'CO_intra', 'OO_intra'\n", 1000 | "\n", 1001 | "for f in ['xtb_md', 'mace01_md']:\n", 1002 | " traj = read('moldyn/'+f+'.xyz', '50:') #ignore first 50 frames\n", 1003 | " for at in traj:\n", 1004 | " at.pbc = True #create a fake box for rdf compatibility\n", 1005 | " at.cell = [100,100,100]\n", 1006 | " rdf = aa.compute_rdfs_traj_avg(traj, rmax=5, nbins=50) #aseMolec provides functionality to compute RDFs\n", 1007 | " plt.plot(rdf[1], rdf[0][tag], '.-', label=f, alpha=0.7, linewidth=3)\n", 1008 | "\n", 1009 | "plt.legend();\n", 1010 | "plt.yticks([]);\n", 1011 | "plt.xlabel(r'R ($\\rm \\AA$)');\n", 1012 | "plt.ylabel('RDF '+tag);" 1013 | ] 1014 | }, 1015 | { 1016 | "cell_type": "markdown", 1017 | "metadata": { 1018 | "cell_id": "04cb8ae4ae644063bbf36dbc565e07b9", 1019 | "deepnote_cell_type": "markdown" 1020 | }, 1021 | "source": [ 1022 | "Try it yourself! Inspect other RDF pairs (`C-O`, `H-C`), how well are they reproduced?\n", 1023 | "\n", 1024 | "The trajectories here are stable and also quite accurate. This is still a relatively simple task: we chose a single molecule at relatively small temperatures (for a molecule) and only ran for 2 picoseconds. In practice, given enough time and high enough temperature the initial models will fail.\n", 1025 | "\n", 1026 | "Experiment with the starting configs, temperatures, simulation length, see if you can find problems with the potentials!" 1027 | ] 1028 | }, 1029 | { 1030 | "cell_type": "markdown", 1031 | "metadata": { 1032 | "cell_id": "6334ee06bcc24a8db4c5a87dfc2d7985", 1033 | "deepnote_cell_type": "markdown", 1034 | "jp-MarkdownHeadingCollapsed": true 1035 | }, 1036 | "source": [ 1037 | "### 4.3 MD of a molecular liquid?" 1038 | ] 1039 | }, 1040 | { 1041 | "cell_type": "markdown", 1042 | "metadata": { 1043 | "cell_id": "bc011009c8a845a7b48ad6f6c4570c39", 1044 | "deepnote_cell_type": "markdown" 1045 | }, 1046 | "source": [ 1047 | "The MLIP was trained on clusters, can we simulate the liquid molecular environment?" 1048 | ] 1049 | }, 1050 | { 1051 | "cell_type": "code", 1052 | "execution_count": null, 1053 | "metadata": { 1054 | "cell_id": "db474b39ce1c4ba1918219a52c749dd1", 1055 | "deepnote_cell_type": "code", 1056 | "deepnote_to_be_reexecuted": false, 1057 | "execution_millis": 495641, 1058 | "execution_start": 1690389440265, 1059 | "source_hash": null 1060 | }, 1061 | "outputs": [], 1062 | "source": [ 1063 | "init_conf = read('data/solvent_liquid.xyz') #read a liquid config with periodic boundary conditions\n", 1064 | "init_conf.center()\n", 1065 | "simpleMD(init_conf, temp=500, calc=mace_calc, fname='moldyn/mace01_md_liquid.xyz', s=10, T=2000)" 1066 | ] 1067 | }, 1068 | { 1069 | "cell_type": "markdown", 1070 | "metadata": { 1071 | "cell_id": "bfa824eb023047b0ab1ce1d1754e1430", 1072 | "deepnote_cell_type": "markdown" 1073 | }, 1074 | "source": [ 1075 | "This XTB calculator is non-periodic, so this dynamics would not be possible without an MLIP! Check for yourself, by replacing the calculator with `xtb`. The system is much larger than the example before (12 molecules vs just one). Let's view the trajectory:" 1076 | ] 1077 | }, 1078 | { 1079 | "cell_type": "code", 1080 | "execution_count": null, 1081 | "metadata": {}, 1082 | "outputs": [], 1083 | "source": [ 1084 | "import nglview as nv\n", 1085 | "traj = read('moldyn/mace01_md_liquid.xyz', ':')\n", 1086 | "view = nv.show_asetraj(traj)\n", 1087 | "view.add_representation('unitcell')\n", 1088 | "view" 1089 | ] 1090 | }, 1091 | { 1092 | "cell_type": "markdown", 1093 | "metadata": { 1094 | "cell_id": "bfa824eb023047b0ab1ce1d1754e1430", 1095 | "deepnote_cell_type": "markdown" 1096 | }, 1097 | "source": [ 1098 | "Transferability from clusters to the condensed phase environment is still an [open research](https://doi.org/10.1021/acs.jpcb.2c03746) question. If this works, it implies that we might be able to learn highly accurate Quantum Chemistry PES on molecular clusters and make predictions (density, diffusivity) for the condensed phase! This is new science!" 1099 | ] 1100 | } 1101 | ], 1102 | "metadata": { 1103 | "deepnote": { 1104 | "is_reactive": false 1105 | }, 1106 | "deepnote_execution_queue": [], 1107 | "deepnote_notebook_id": "2a16e1e8a7ad42a8825007f5cff15d19", 1108 | "kernelspec": { 1109 | "display_name": "Python 3 (ipykernel)", 1110 | "language": "python", 1111 | "name": "python3" 1112 | }, 1113 | "language_info": { 1114 | "codemirror_mode": { 1115 | "name": "ipython", 1116 | "version": 3 1117 | }, 1118 | "file_extension": ".py", 1119 | "mimetype": "text/x-python", 1120 | "name": "python", 1121 | "nbconvert_exporter": "python", 1122 | "pygments_lexer": "ipython3", 1123 | "version": "3.12.4" 1124 | } 1125 | }, 1126 | "nbformat": 4, 1127 | "nbformat_minor": 4 1128 | } 1129 | -------------------------------------------------------------------------------- /T02-MACE-Practice-II.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | " !pip install xtb mace-torch nglview ipywidgets\n", 10 | " !pip install git+https://github.com/imagdau/aseMolec@main" 11 | ] 12 | }, 13 | { 14 | "cell_type": "markdown", 15 | "metadata": { 16 | "cell_id": "a40121600b3c467da428477bc8bf55e0", 17 | "deepnote_cell_type": "markdown" 18 | }, 19 | "source": [ 20 | "# MACE in Practice II" 21 | ] 22 | }, 23 | { 24 | "cell_type": "markdown", 25 | "metadata": { 26 | "cell_id": "95057ab91ce54145b569a1b35cc81742", 27 | "deepnote_cell_type": "markdown", 28 | "tags": [] 29 | }, 30 | "source": [ 31 | "In this tutorial, you will learn how to improve MLIP models by using iterative training and active learning. We illustrate these training workflows on MACE, but they are broadly applicable to all MLIPs. We will also showcase the state-of-the-art [foundational models](https://matbench-discovery.materialsproject.org/) - the latest development in the field of MLIPs. These models are trained on massive training sets of [inorganic](https://doi.org/10.48550/arXiv.2401.00096) and [organic](https://doi.org/10.48550/arXiv.2312.15211) databases and show a great deal of `out-of-the-box` MD stability in an extensive variety of [applications](https://doi.org/10.48550/arXiv.2401.00096). We will discuss [fine-tunning](https://doi.org/10.48550/arXiv.2405.20217) which is an actively-researched technique to tweak these foundational models to new systems (out of training) and/or new levels of reference theory.\n", 32 | "\n", 33 | "## Learning Objectives for today:\n", 34 | "\n", 35 | "1. **Iterative Training: improving stability and accuracy**\n", 36 | "2. **Active learning: committee models**\n", 37 | "3. **Foundational models and fine-tuning**" 38 | ] 39 | }, 40 | { 41 | "cell_type": "markdown", 42 | "metadata": { 43 | "cell_id": "6334ee06bcc24a8db4c5a87dfc2d7985", 44 | "deepnote_cell_type": "markdown" 45 | }, 46 | "source": [ 47 | "## 1. Iterative Training" 48 | ] 49 | }, 50 | { 51 | "cell_type": "markdown", 52 | "metadata": {}, 53 | "source": [ 54 | "### 1.1 MD with a smaller MACE model" 55 | ] 56 | }, 57 | { 58 | "cell_type": "markdown", 59 | "metadata": { 60 | "cell_id": "bfa824eb023047b0ab1ce1d1754e1430", 61 | "deepnote_cell_type": "markdown" 62 | }, 63 | "source": [ 64 | "The model we trained in our previous tutorial was already stable in MD and quite accurate with little training. This is both because MACE models are smooth and regular (combining lessons from 10-15 years of MLIP development) but also because the task of simulating a single molecule for a few picoseconds is not all that difficult. In general, in real research applications, achieving MD stability and accuracy is not always straightforward from the get-go. Models can be improved through iterative training and active learning which expands the training data to fix errors on the model's potential energy surface. To illustrate these concepts in practice, let's first train a less accurate MACE by reducing the model size and amount of training data:" 65 | ] 66 | }, 67 | { 68 | "cell_type": "code", 69 | "execution_count": null, 70 | "metadata": {}, 71 | "outputs": [], 72 | "source": [ 73 | "from ase.io import read, write\n", 74 | "db = read('data/solvent_xtb.xyz', ':')\n", 75 | "write('data/solvent_xtb_train_50.xyz', db[:53]) #first 50 configs plus the 3 E0s, we'll use only 20 of them for training!\n", 76 | "\n", 77 | "!mace_run_train \\\n", 78 | " --model=\"MACE\" \\\n", 79 | " --num_interactions=2 \\\n", 80 | " --num_channels=16 \\\n", 81 | " --max_L=0 \\\n", 82 | " --correlation=2 \\\n", 83 | " --r_max=4.0 \\\n", 84 | " --max_ell=2 \\\n", 85 | " --name=\"mace02_com1\" \\\n", 86 | " --model_dir=\"MACE_models\" \\\n", 87 | " --log_dir=\"MACE_models\" \\\n", 88 | " --checkpoints_dir=\"MACE_models\" \\\n", 89 | " --results_dir=\"MACE_models\" \\\n", 90 | " --train_file=\"data/solvent_xtb_train_20.xyz\" \\\n", 91 | " --valid_file=\"data/solvent_xtb_train_50.xyz\" \\\n", 92 | " --test_file=\"data/solvent_xtb_test.xyz\" \\\n", 93 | " --E0s=\"average\" \\\n", 94 | " --energy_key=\"energy_xtb\" \\\n", 95 | " --forces_key=\"forces_xtb\" \\\n", 96 | " --device=cuda \\\n", 97 | " --batch_size=10 \\\n", 98 | " --max_num_epochs=300 \\\n", 99 | " --swa \\\n", 100 | " --seed=123\n", 101 | "\n", 102 | "#remove checkpoints since they may cause errors on retraining a model with the same name but a different architecture\n", 103 | "import glob\n", 104 | "import os\n", 105 | "for file in glob.glob(\"MACE_models/*_run-*.model\"):\n", 106 | " os.remove(file)\n", 107 | "for file in glob.glob(\"MACE_models/*.pt\"):\n", 108 | " os.remove(file)" 109 | ] 110 | }, 111 | { 112 | "cell_type": "markdown", 113 | "metadata": {}, 114 | "source": [ 115 | "Notice, we are getting substantially larger errors than before. Now, let's run some dynamics:" 116 | ] 117 | }, 118 | { 119 | "cell_type": "code", 120 | "execution_count": null, 121 | "metadata": {}, 122 | "outputs": [], 123 | "source": [ 124 | "from ase.io import read, write\n", 125 | "from ase import units\n", 126 | "from ase.md.langevin import Langevin\n", 127 | "from ase.md.velocitydistribution import Stationary, ZeroRotation, MaxwellBoltzmannDistribution\n", 128 | "from aseMolec import extAtoms as ea\n", 129 | "\n", 130 | "import os\n", 131 | "import time\n", 132 | "import numpy as np\n", 133 | "import pylab as pl\n", 134 | "from IPython import display\n", 135 | "np.random.seed(701) #just making sure the MD failure is reproducible\n", 136 | "\n", 137 | "def simpleMD(init_conf, temp, calc, fname, s, T):\n", 138 | " init_conf.set_calculator(calc)\n", 139 | "\n", 140 | " #initialize the temperature\n", 141 | " \n", 142 | " MaxwellBoltzmannDistribution(init_conf, temperature_K=300) #initialize temperature at 300\n", 143 | " Stationary(init_conf)\n", 144 | " ZeroRotation(init_conf)\n", 145 | "\n", 146 | " dyn = Langevin(init_conf, 1.0*units.fs, temperature_K=temp, friction=0.1) #drive system to desired temperature\n", 147 | "\n", 148 | " %matplotlib inline\n", 149 | "\n", 150 | " time_fs = []\n", 151 | " temperature = []\n", 152 | " energies = []\n", 153 | "\n", 154 | " #remove previously stored trajectory with the same name\n", 155 | " os.system('rm -rfv '+fname)\n", 156 | "\n", 157 | " fig, ax = pl.subplots(2, 1, figsize=(6,6), sharex='all', gridspec_kw={'hspace': 0, 'wspace': 0})\n", 158 | "\n", 159 | " def write_frame():\n", 160 | " dyn.atoms.info['energy_mace'] = dyn.atoms.get_potential_energy()\n", 161 | " dyn.atoms.arrays['force_mace'] = dyn.atoms.calc.get_forces()\n", 162 | " dyn.atoms.write(fname, append=True)\n", 163 | " time_fs.append(dyn.get_time()/units.fs)\n", 164 | " temperature.append(dyn.atoms.get_temperature())\n", 165 | " energies.append(dyn.atoms.get_potential_energy()/len(dyn.atoms))\n", 166 | "\n", 167 | " ax[0].plot(np.array(time_fs), np.array(energies), color=\"b\")\n", 168 | " ax[0].set_ylabel('E (eV/atom)')\n", 169 | "\n", 170 | " # plot the temperature of the system as subplots\n", 171 | " ax[1].plot(np.array(time_fs), temperature, color=\"r\")\n", 172 | " ax[1].set_ylabel('T (K)')\n", 173 | " ax[1].set_xlabel('Time (fs)')\n", 174 | "\n", 175 | " display.clear_output(wait=True)\n", 176 | " display.display(pl.gcf())\n", 177 | " time.sleep(0.01)\n", 178 | "\n", 179 | " dyn.attach(write_frame, interval=s)\n", 180 | " t0 = time.time()\n", 181 | " dyn.run(T)\n", 182 | " t1 = time.time()\n", 183 | " print(\"MD finished in {0:.2f} minutes!\".format((t1-t0)/60))\n", 184 | "\n", 185 | "#let us start with a single molecule\n", 186 | "init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n", 187 | "\n", 188 | "#we can use MACE as a calculator in ASE!\n", 189 | "from mace.calculators import MACECalculator\n", 190 | "mace_calc = MACECalculator(model_paths=['MACE_models/mace02_com1_swa_compiled.model'], device='cuda')\n", 191 | "\n", 192 | "simpleMD(init_conf, temp=1200, calc=mace_calc, fname='moldyn/mace02_md.xyz', s=10, T=2000)" 193 | ] 194 | }, 195 | { 196 | "cell_type": "markdown", 197 | "metadata": {}, 198 | "source": [ 199 | "Depending on the random seed for the MD, you may see different things. At first sight, the energy against time doesn't look too bad, although the long time scale wandering of the energy is a little weird. Lets look at the trajectory." 200 | ] 201 | }, 202 | { 203 | "cell_type": "code", 204 | "execution_count": null, 205 | "metadata": {}, 206 | "outputs": [], 207 | "source": [ 208 | "from ase.io import read, write\n", 209 | "import nglview as nv\n", 210 | "\n", 211 | "traj = read('moldyn/mace02_md.xyz', ':')\n", 212 | "view = nv.show_asetraj(traj)\n", 213 | "view._set_size(w='100%', h='500px')\n", 214 | "view" 215 | ] 216 | }, 217 | { 218 | "cell_type": "markdown", 219 | "metadata": {}, 220 | "source": [ 221 | "If you go to the end of the trajectory, you should find that the bond angles are actually very strange - it looks unphsyical." 222 | ] 223 | }, 224 | { 225 | "cell_type": "markdown", 226 | "metadata": {}, 227 | "source": [ 228 | "### 1.2 Identify the problem and expand training\n", 229 | "\n", 230 | "Something doesn't look right with the hydrogen atoms. Let's re-evaluate the first 100 steps from the trajectory on the reference XTB potential energy surface and then plot it against MACE energy." 231 | ] 232 | }, 233 | { 234 | "cell_type": "code", 235 | "execution_count": null, 236 | "metadata": {}, 237 | "outputs": [], 238 | "source": [ 239 | "from ase.io import read, write\n", 240 | "from tqdm import tqdm\n", 241 | "from xtb.ase.calculator import XTB\n", 242 | "xtb_calc = XTB(method=\"GFN2-xTB\")\n", 243 | "\n", 244 | "#compute true reference XTB values\n", 245 | "traj = read('moldyn/mace02_md.xyz', ':')\n", 246 | "for at in tqdm(traj[:100]):\n", 247 | " at.calc = None\n", 248 | " at.calc = xtb_calc\n", 249 | " at.info['energy_xtb'] = at.get_potential_energy()\n", 250 | " at.arrays['forces_xtb'] = at.get_forces()\n", 251 | " at.calc = None\n", 252 | "\n", 253 | "write('data/mace02_md_100_xtb.xyz', traj[:100]) #save full result" 254 | ] 255 | }, 256 | { 257 | "cell_type": "code", 258 | "execution_count": null, 259 | "metadata": {}, 260 | "outputs": [], 261 | "source": [ 262 | "import numpy as np\n", 263 | "from aseMolec import extAtoms as ea\n", 264 | "from matplotlib import pyplot as plt\n", 265 | "\n", 266 | "traj = read('data/mace02_md_100_xtb.xyz', ':')\n", 267 | "plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_xtb', peratom=True), label='XTB')\n", 268 | "plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace', peratom=True), label='MACE')\n", 269 | "plt.legend()\n", 270 | "plt.xlabel('Time (fs)')\n", 271 | "plt.ylabel('Total Energy per Atom (eV)')" 272 | ] 273 | }, 274 | { 275 | "cell_type": "markdown", 276 | "metadata": {}, 277 | "source": [ 278 | "Indeed, at around frame 40 the XTB energy diverges, because MACE finds an unphysical config. With older potentials, at this point the MD would explode, but because MACE is much smoother the simulation keeps going, albeit generating the wrong dynamics. Let's take three of these failed configs, add them back to the training set and refit a new model. This is called iterative training:\n", 279 | "![alt text](figures/iterative_training.png)" 280 | ] 281 | }, 282 | { 283 | "cell_type": "code", 284 | "execution_count": null, 285 | "metadata": {}, 286 | "outputs": [], 287 | "source": [ 288 | "db = read('data/solvent_xtb_train_20.xyz', ':')\n", 289 | "db += traj[40:100:20] #add three failed configs to the training set\n", 290 | "write('data/solvent_xtb_train_23_gen1.xyz', db)" 291 | ] 292 | }, 293 | { 294 | "cell_type": "markdown", 295 | "metadata": {}, 296 | "source": [ 297 | "### 1.3 Train a new MACE model and run MD again" 298 | ] 299 | }, 300 | { 301 | "cell_type": "code", 302 | "execution_count": null, 303 | "metadata": {}, 304 | "outputs": [], 305 | "source": [ 306 | "!mace_run_train \\\n", 307 | " --model=\"MACE\" \\\n", 308 | " --num_interactions=2 \\\n", 309 | " --num_channels=16 \\\n", 310 | " --max_L=0 \\\n", 311 | " --correlation=2 \\\n", 312 | " --r_max=4.0 \\\n", 313 | " --max_ell=2 \\\n", 314 | " --name=\"mace02_com1_gen1\" \\\n", 315 | " --model_dir=\"MACE_models\" \\\n", 316 | " --log_dir=\"MACE_models\" \\\n", 317 | " --checkpoints_dir=\"MACE_models\" \\\n", 318 | " --results_dir=\"MACE_models\" \\\n", 319 | " --train_file=\"data/solvent_xtb_train_23_gen1.xyz\" \\\n", 320 | " --valid_file=\"data/solvent_xtb_train_50.xyz\" \\\n", 321 | " --test_file=\"data/solvent_xtb_test.xyz\" \\\n", 322 | " --E0s=\"average\" \\\n", 323 | " --energy_key=\"energy_xtb\" \\\n", 324 | " --forces_key=\"forces_xtb\" \\\n", 325 | " --device=cuda \\\n", 326 | " --batch_size=10 \\\n", 327 | " --max_num_epochs=300 \\\n", 328 | " --swa \\\n", 329 | " --seed=123\n", 330 | "\n", 331 | "#remove checkpoints since they may cause errors on retraining a model with the same name but a different architecture\n", 332 | "import glob\n", 333 | "import os\n", 334 | "for file in glob.glob(\"MACE_models/*_run-*.model\"):\n", 335 | " os.remove(file)\n", 336 | "for file in glob.glob(\"MACE_models/*.pt\"):\n", 337 | " os.remove(file)" 338 | ] 339 | }, 340 | { 341 | "cell_type": "markdown", 342 | "metadata": {}, 343 | "source": [ 344 | "For some reason the energy error on the training set is now huge - can you work out why this is?\n", 345 | "\n", 346 | "What does this imply about how we do iterative training?" 347 | ] 348 | }, 349 | { 350 | "cell_type": "code", 351 | "execution_count": null, 352 | "metadata": {}, 353 | "outputs": [], 354 | "source": [ 355 | "mace_calc = MACECalculator(model_paths=['MACE_models/mace02_com1_gen1_swa_compiled.model'], device='cuda')\n", 356 | "init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n", 357 | "simpleMD(init_conf, temp=1200, calc=mace_calc, fname='moldyn/mace02_md_gen1.xyz', s=10, T=2000)" 358 | ] 359 | }, 360 | { 361 | "cell_type": "code", 362 | "execution_count": null, 363 | "metadata": {}, 364 | "outputs": [], 365 | "source": [ 366 | "traj = read('moldyn/mace02_md_gen1.xyz', ':')\n", 367 | "view = nv.show_asetraj(traj)\n", 368 | "view._set_size(w='100%', h='500px')\n", 369 | "view" 370 | ] 371 | }, 372 | { 373 | "cell_type": "markdown", 374 | "metadata": {}, 375 | "source": [ 376 | "Great! The dynamics is already looking better, however its hard to tell if it is really correct. To do this we cuold look at the radial distribution function compared to a ground truth trajectory, but if the ground truth is too expensive its not so easy.\n", 377 | "\n", 378 | "If we have reason to believe the model is wrong, we could continue the iterative process and gradually improve the potential. This is an arduous process, because we need to carefully investigate the trajectories and decide which configs to add back to training. We could instead automate this protocol by predicting errors on the fly and picking configs which are not well predicted: this is called active learning." 379 | ] 380 | }, 381 | { 382 | "cell_type": "markdown", 383 | "metadata": { 384 | "cell_id": "a00c92d61afc4adcb493f1cc7b2ee3a6", 385 | "deepnote_cell_type": "markdown" 386 | }, 387 | "source": [ 388 | "## 2. Active Learning with MACE" 389 | ] 390 | }, 391 | { 392 | "cell_type": "markdown", 393 | "metadata": {}, 394 | "source": [ 395 | "### 2.1 Preparing a committee of models" 396 | ] 397 | }, 398 | { 399 | "cell_type": "markdown", 400 | "metadata": { 401 | "cell_id": "588eeeda5dc94c76841aba2af1c4f90a", 402 | "deepnote_cell_type": "markdown" 403 | }, 404 | "source": [ 405 | "We can compute errors by evaluating the reference energy and forces (in our case XTB) and computing the difference to MACE predictions. In real research applications, this can be very expensive to evaluate depending on the referneece level of theory. Alternatively, we can estimate errors based on a committee of models. Let's train a committee of MACE models by adding some randomness to the optimization process. We can achieve this by changing the `--seed`. We already have a model, we will fit two more, on the same data:" 406 | ] 407 | }, 408 | { 409 | "cell_type": "code", 410 | "execution_count": null, 411 | "metadata": {}, 412 | "outputs": [], 413 | "source": [ 414 | "!mace_run_train \\\n", 415 | " --model=\"MACE\" \\\n", 416 | " --num_interactions=2 \\\n", 417 | " --num_channels=16 \\\n", 418 | " --max_L=0 \\\n", 419 | " --correlation=2 \\\n", 420 | " --r_max=4.0 \\\n", 421 | " --max_ell=2 \\\n", 422 | " --name=\"mace02_com2\" \\\n", 423 | " --model_dir=\"MACE_models\" \\\n", 424 | " --log_dir=\"MACE_models\" \\\n", 425 | " --checkpoints_dir=\"MACE_models\" \\\n", 426 | " --results_dir=\"MACE_models\" \\\n", 427 | " --train_file=\"data/solvent_xtb_train_50.xyz\" \\\n", 428 | " --valid_fraction=0.60 \\\n", 429 | " --test_file=\"data/solvent_xtb_test.xyz\" \\\n", 430 | " --E0s=\"average\" \\\n", 431 | " --energy_key=\"energy_xtb\" \\\n", 432 | " --forces_key=\"forces_xtb\" \\\n", 433 | " --device=cuda \\\n", 434 | " --batch_size=10 \\\n", 435 | " --max_num_epochs=300 \\\n", 436 | " --swa \\\n", 437 | " --seed=345\n", 438 | "\n", 439 | "!mace_run_train \\\n", 440 | " --model=\"MACE\" \\\n", 441 | " --num_interactions=2 \\\n", 442 | " --num_channels=16 \\\n", 443 | " --max_L=0 \\\n", 444 | " --correlation=2 \\\n", 445 | " --r_max=4.0 \\\n", 446 | " --max_ell=2 \\\n", 447 | " --name=\"mace02_com3\" \\\n", 448 | " --model_dir=\"MACE_models\" \\\n", 449 | " --log_dir=\"MACE_models\" \\\n", 450 | " --checkpoints_dir=\"MACE_models\" \\\n", 451 | " --results_dir=\"MACE_models\" \\\n", 452 | " --train_file=\"data/solvent_xtb_train_50.xyz\" \\\n", 453 | " --valid_fraction=0.60 \\\n", 454 | " --test_file=\"data/solvent_xtb_test.xyz\" \\\n", 455 | " --E0s=\"average\" \\\n", 456 | " --energy_key=\"energy_xtb\" \\\n", 457 | " --forces_key=\"forces_xtb\" \\\n", 458 | " --device=cuda \\\n", 459 | " --batch_size=10 \\\n", 460 | " --max_num_epochs=300 \\\n", 461 | " --swa \\\n", 462 | " --seed=567\n", 463 | "\n", 464 | "#remove checkpoints since they may cause errors on retraining a model with the same name but a different architecture\n", 465 | "import glob\n", 466 | "import os\n", 467 | "for file in glob.glob(\"MACE_models/*_run-*.model\"):\n", 468 | " os.remove(file)\n", 469 | "for file in glob.glob(\"MACE_models/*.pt\"):\n", 470 | " os.remove(file)" 471 | ] 472 | }, 473 | { 474 | "cell_type": "markdown", 475 | "metadata": {}, 476 | "source": [ 477 | "Perfect, we have two new models. Let's start by testing the commitee on the first 100 frames of the first trajectory we generated. The `MACECalculator` can conveniently take a list of calculators as input and will compute separate energies from each calculator. " 478 | ] 479 | }, 480 | { 481 | "cell_type": "code", 482 | "execution_count": null, 483 | "metadata": {}, 484 | "outputs": [], 485 | "source": [ 486 | "import numpy as np\n", 487 | "from ase.io import read\n", 488 | "from aseMolec import extAtoms as ea\n", 489 | "from matplotlib import pyplot as plt\n", 490 | "from mace.calculators import MACECalculator\n", 491 | "\n", 492 | "model_paths = ['MACE_models/mace02_com1_swa_compiled.model',\n", 493 | " 'MACE_models/mace02_com2_swa_compiled.model',\n", 494 | " 'MACE_models/mace02_com3_swa_compiled.model']\n", 495 | "mace_calcs = MACECalculator(model_paths=model_paths, device='cpu')\n", 496 | "\n", 497 | "traj = read('data/mace02_md_100_xtb.xyz', ':')\n", 498 | "for at in traj:\n", 499 | " at.calc = mace_calcs\n", 500 | " engs = at.get_potential_energies()\n", 501 | " at.info['energy_mace_1'] = at.info.pop('energy_mace') #rename value obtained with first member of the committee\n", 502 | " at.info['energy_mace_2'] = engs[1]\n", 503 | " at.info['energy_mace_3'] = engs[2]\n", 504 | " at.info['variance'] = np.std(engs)\n", 505 | " at.info['average_mace_energy'] = np.average(engs)\n", 506 | " at.info['true_error'] = np.abs(at.info['average_mace_energy'] - at.info['energy_xtb'])\n", 507 | " \n", 508 | "#Let's check the energies of the MACE committee vs XTB energy\n", 509 | "plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_xtb', peratom=True), label='XTB')\n", 510 | "plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace_1', peratom=True), label='MACE_1')\n", 511 | "plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace_2', peratom=True), label='MACE_2')\n", 512 | "plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace_3', peratom=True), label='MACE_3')\n", 513 | "plt.legend()\n", 514 | "plt.xlabel('Time (fs)')\n", 515 | "plt.ylabel('Energy per Atom (eV)')\n", 516 | "plt.show()\n", 517 | "\n", 518 | "\n", 519 | "#plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'variance', peratom=True), label='committee variance')\n", 520 | "#plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'true_error', peratom=True), label='error w.r.t XTB')\n", 521 | "\n", 522 | "fig, ax1 = plt.subplots()\n", 523 | "ax2 = ax1.twinx()\n", 524 | "ax1.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'variance', peratom=True), label='committee variance', color='tab:blue')\n", 525 | "#ax2.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'true_error', peratom=True), label='error w.r.t XTB', color='tab:orange')\n", 526 | "\n", 527 | "ax1.set_xlabel('time (fs)')\n", 528 | "ax1.set_ylabel('committee energy variance', color='tab:blue')\n", 529 | "#ax2.set_ylabel('error w.r.t XTB', color='tab:orange')\n", 530 | "\n", 531 | "plt.show()" 532 | ] 533 | }, 534 | { 535 | "cell_type": "markdown", 536 | "metadata": { 537 | "cell_id": "ec2e91aa94044d14aeadea824b5fecbe", 538 | "deepnote_cell_type": "markdown" 539 | }, 540 | "source": [ 541 | "Notice how the variance (disagreement between models) increases around the same config where the true error with respect to XTB diverges. This is good news because it indicates the variance is a good proxy for true error.\n", 542 | "\n", 543 | "Now we can run dynamics with a commitee of models and monitor the variance in the energy prediction. Because XTB is cheap enough we can also compare that variance with the true error. Do they correlate?" 544 | ] 545 | }, 546 | { 547 | "cell_type": "markdown", 548 | "metadata": {}, 549 | "source": [ 550 | "### 2.2 Running MD with the MACE committee" 551 | ] 552 | }, 553 | { 554 | "cell_type": "code", 555 | "execution_count": null, 556 | "metadata": { 557 | "cell_id": "763720f0e2ed4ea7a27db276507b4fa9", 558 | "deepnote_cell_type": "code", 559 | "deepnote_to_be_reexecuted": false, 560 | "execution_millis": 741598, 561 | "execution_start": 1690390325607, 562 | "source_hash": null 563 | }, 564 | "outputs": [], 565 | "source": [ 566 | "from aseMolec import extAtoms as ea\n", 567 | "from ase import units\n", 568 | "from ase.md.langevin import Langevin\n", 569 | "from ase.md.velocitydistribution import Stationary, ZeroRotation, MaxwellBoltzmannDistribution\n", 570 | "from ase.io import read, write\n", 571 | "\n", 572 | "import random\n", 573 | "import numpy as np\n", 574 | "import time\n", 575 | "import pylab as pl\n", 576 | "from IPython import display\n", 577 | "\n", 578 | "from mace.calculators import MACECalculator\n", 579 | "from xtb.ase.calculator import XTB\n", 580 | "\n", 581 | "model_paths = ['MACE_models/mace02_com1_swa_compiled.model',\n", 582 | " 'MACE_models/mace02_com2_swa_compiled.model',\n", 583 | " 'MACE_models/mace02_com3_swa_compiled.model']\n", 584 | "xtb_calc = XTB(method=\"GFN2-xTB\")\n", 585 | "mace_calc = MACECalculator(model_paths=model_paths, device='cpu')\n", 586 | "\n", 587 | "init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n", 588 | "init_conf.calc = mace_calc\n", 589 | "\n", 590 | "#initialize the temperature\n", 591 | "np.random.seed(701)\n", 592 | "MaxwellBoltzmannDistribution(init_conf, temperature_K=300)\n", 593 | "Stationary(init_conf)\n", 594 | "ZeroRotation(init_conf)\n", 595 | "\n", 596 | "dyn = Langevin(init_conf, 1*units.fs, temperature_K=1200, friction=0.1)\n", 597 | "\n", 598 | "%matplotlib inline\n", 599 | "\n", 600 | "time_fs = []\n", 601 | "temperature = []\n", 602 | "energies_1 = []\n", 603 | "energies_2 = []\n", 604 | "energies_3 = []\n", 605 | "variances = []\n", 606 | "xtb_energies = []\n", 607 | "true_errors = []\n", 608 | "\n", 609 | "!rm -rfv moldyn/mace02_md_committee.xyz\n", 610 | "fig, ax = pl.subplots(4, 1, figsize=(8,8), sharex='all', gridspec_kw={'hspace': 0, 'wspace': 0})\n", 611 | "\n", 612 | "\n", 613 | "def write_frame():\n", 614 | " at = dyn.atoms.copy()\n", 615 | " at.calc = xtb_calc\n", 616 | " xtb_energy = at.get_potential_energy()\n", 617 | "\n", 618 | " dyn.atoms.write('moldyn/mace02_md_committee.xyz', append=True, write_results=False)\n", 619 | " time_fs.append(dyn.get_time()/units.fs)\n", 620 | " temperature.append(dyn.atoms.get_temperature())\n", 621 | " energies_1.append(dyn.atoms.calc.results[\"energies\"][0]/len(dyn.atoms))\n", 622 | " energies_2.append(dyn.atoms.calc.results[\"energies\"][1]/len(dyn.atoms))\n", 623 | " energies_3.append(dyn.atoms.calc.results[\"energies\"][2]/len(dyn.atoms))\n", 624 | " variances.append(dyn.atoms.calc.results[\"energy_var\"]/len(dyn.atoms))\n", 625 | " xtb_energies.append(xtb_energy/len(dyn.atoms))\n", 626 | " true_errors.append(np.var([dyn.atoms.calc.results[\"energy\"],xtb_energy])/len(dyn.atoms))\n", 627 | "\n", 628 | " # plot the true error\n", 629 | " ax[0].plot(np.array(time_fs), np.array(true_errors), color=\"black\")\n", 630 | " ax[0].set_ylabel(r'$\\Delta$ E (eV$^2$/atom)')\n", 631 | " ax[0].legend(['Error w.r.t. XTB'], loc='upper left')\n", 632 | "\n", 633 | " # plot committee variance\n", 634 | " ax[1].plot(np.array(time_fs), np.array(variances), color=\"y\")\n", 635 | " ax[1].set_ylabel(r'committee variance')\n", 636 | " ax[1].legend(['Estimated Error (committee variances)'], loc='upper left')\n", 637 | "\n", 638 | " # plot the temperature of the system as subplots\n", 639 | " ax[2].plot(np.array(time_fs), temperature, color=\"r\", label='Temperature')\n", 640 | " ax[2].set_ylabel(\"T (K)\")\n", 641 | "\n", 642 | " ax[3].plot(np.array(time_fs), energies_1, color=\"g\")\n", 643 | " ax[3].plot(np.array(time_fs), energies_2, color=\"y\")\n", 644 | " ax[3].plot(np.array(time_fs), energies_3, color=\"olive\")\n", 645 | " ax[3].plot(np.array(time_fs), xtb_energies, color=\"black\")\n", 646 | " ax[3].set_ylabel(\"E (eV/atom)\")\n", 647 | " ax[3].set_xlabel('Time (fs)')\n", 648 | " ax[3].legend(['E mace1', 'E mace2', 'E mace3', 'E xtb'], loc='upper left')\n", 649 | "\n", 650 | " display.clear_output(wait=True)\n", 651 | " display.display(fig)\n", 652 | " time.sleep(0.01)\n", 653 | "\n", 654 | "dyn.attach(write_frame, interval=10)\n", 655 | "dyn.run(2000)\n", 656 | "print(\"MD finished!\")" 657 | ] 658 | }, 659 | { 660 | "cell_type": "markdown", 661 | "metadata": {}, 662 | "source": [ 663 | "NOTE: if you get the error `xtb could not evalute the config` that means the dynamics went so crazy and gave such strange configurations, that xtb refused to run! thats to be expected if the model is really bad. Copy some of the code above to have a look at the trajectory if you'd like to see this. \n", 664 | "\n", 665 | "Closely observe the dynamics. Notice how good the committee error is as a proxy for the true error. In this case the true is cheap to compute, but in most practical applications it won't be. Therefore, we will need to rely on the committee error to identy configurations that should be added back to the training set. This is called active learning:" 666 | ] 667 | }, 668 | { 669 | "cell_type": "markdown", 670 | "metadata": {}, 671 | "source": [ 672 | "![alt text](figures/active_learning.png)" 673 | ] 674 | }, 675 | { 676 | "cell_type": "markdown", 677 | "metadata": {}, 678 | "source": [ 679 | "### Active learning in practice\n", 680 | "\n", 681 | "The way to use active learning to improve the model is as follows:\n", 682 | "1. run dynmics, track the uncertainty.\n", 683 | "2. if the uncertainty breaches some predeterined value, stop the simulation and peform the ground truth calculation.\n", 684 | "3. and the new config to the dataset, and retrain\n", 685 | "4. repeat steps 1-3 until the uncertainty never crosses the threshold\n", 686 | "\n", 687 | "This can be done without human supervision - you can write a program which loops this process. \n", 688 | "\n", 689 | "As an exercise at the end of the notebook, try writing an active learing loop to gradually grow the dataset and produce a good model, without ever running XTB dynamics. " 690 | ] 691 | }, 692 | { 693 | "cell_type": "markdown", 694 | "metadata": {}, 695 | "source": [ 696 | "## 3 Foundational Models" 697 | ] 698 | }, 699 | { 700 | "cell_type": "markdown", 701 | "metadata": {}, 702 | "source": [ 703 | "### 3.1 Molecular Dynamics with MACE-MP-0" 704 | ] 705 | }, 706 | { 707 | "cell_type": "markdown", 708 | "metadata": {}, 709 | "source": [ 710 | "Foundation models changed everything. MACE-MP-0 is a model trained on >1 million DFT calcuations, and can run dynamics for the whole periodic table. \n", 711 | "\n", 712 | "Mace provides a simple interface to load a foundational model, which we can use mow. Check the [documentation](https://mace-docs.readthedocs.io/en/latest/guide/foundation_models.html) for more details." 713 | ] 714 | }, 715 | { 716 | "cell_type": "code", 717 | "execution_count": null, 718 | "metadata": {}, 719 | "outputs": [], 720 | "source": [ 721 | "from mace.calculators import mace_mp\n", 722 | "\n", 723 | "macemp = mace_mp(model=\"small\")\n", 724 | "init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n", 725 | "simpleMD(init_conf, temp=1200, calc=macemp, fname='moldyn/mace03_md.xyz', s=10, T=2000)" 726 | ] 727 | }, 728 | { 729 | "cell_type": "markdown", 730 | "metadata": {}, 731 | "source": [ 732 | "OMG, look at that dynamics! Stable out-of-the-box, sign me up! Let's view the trajectory:" 733 | ] 734 | }, 735 | { 736 | "cell_type": "code", 737 | "execution_count": null, 738 | "metadata": {}, 739 | "outputs": [], 740 | "source": [ 741 | "from ase.io import read, write\n", 742 | "import nglview as nv\n", 743 | "\n", 744 | "traj = read('moldyn/mace03_md.xyz', ':')\n", 745 | "view = nv.show_asetraj(traj)\n", 746 | "view._set_size(w='100%', h='500px')\n", 747 | "view" 748 | ] 749 | }, 750 | { 751 | "cell_type": "markdown", 752 | "metadata": {}, 753 | "source": [ 754 | "### 3.2 Compare to XTB" 755 | ] 756 | }, 757 | { 758 | "cell_type": "markdown", 759 | "metadata": {}, 760 | "source": [ 761 | "Let's compute the radial distribution functions of this very stable trajectory and compare them to XTB. Remember MACE-MP was trained on PBE level of theory so we don't necessarily expect them to match:" 762 | ] 763 | }, 764 | { 765 | "cell_type": "code", 766 | "execution_count": null, 767 | "metadata": {}, 768 | "outputs": [], 769 | "source": [ 770 | "from matplotlib import pyplot as plt\n", 771 | "from aseMolec import anaAtoms as aa\n", 772 | "\n", 773 | "tag = 'OO_intra' #choose one of 'HH_intra', 'HC_intra', 'HO_intra', 'CC_intra', 'CO_intra', 'OO_intra'\n", 774 | "\n", 775 | "for f in ['xtb_md', 'mace03_md']:\n", 776 | " traj = read('moldyn/'+f+'.xyz', '50:') #ignore first 50 frames\n", 777 | " for at in traj:\n", 778 | " at.pbc = True #create a fake box for rdf compatibility\n", 779 | " at.cell = [100,100,100]\n", 780 | " rdf = aa.compute_rdfs_traj_avg(traj, rmax=5, nbins=70) #aseMolec provides functionality to compute RDFs\n", 781 | " plt.plot(rdf[1], rdf[0][tag], '.-', label=f, alpha=0.7, linewidth=3)\n", 782 | "\n", 783 | "plt.legend();\n", 784 | "plt.yticks([]);\n", 785 | "plt.xlabel(r'R ($\\rm \\AA$)');\n", 786 | "plt.ylabel('RDF '+tag);" 787 | ] 788 | }, 789 | { 790 | "cell_type": "markdown", 791 | "metadata": {}, 792 | "source": [ 793 | "Notice there's a slight shift in the O-O RDF peak. This is likely due to the different level of reference theory. Can we fix by fine tuning MACE-MP-0?\n", 794 | "\n", 795 | "Depending on your application, the PBE reference using in MACE-MP-0 may not be appropriate. It is probably better than XTB, but thats because for this notebook we are using XTB since it is extremely fast, so we can check the model's true error very easily. \n", 796 | "\n", 797 | "In practice, you might want to run dynamics of this small molecule at the MP2 or coupled cluster level (which is the 'gold standard' of molecular quantum chemistry). In that case, you would want to finetune MACE-MP-0 onto a small amount of very expensive couple cluster data." 798 | ] 799 | }, 800 | { 801 | "cell_type": "markdown", 802 | "metadata": {}, 803 | "source": [ 804 | "### 3.3 Fine tune MACE-MP to XTB" 805 | ] 806 | }, 807 | { 808 | "cell_type": "code", 809 | "execution_count": null, 810 | "metadata": {}, 811 | "outputs": [], 812 | "source": [ 813 | "!mace_run_train \\\n", 814 | " --name=\"finetuned_MACE\" \\\n", 815 | " --foundation_model=\"small\" \\\n", 816 | " --train_file=\"data/solvent_xtb_train_50.xyz\" \\\n", 817 | " --valid_fraction=0.60 \\\n", 818 | " --test_file=\"data/solvent_xtb_test.xyz\" \\\n", 819 | " --energy_weight=1.0 \\\n", 820 | " --forces_weight=1.0 \\\n", 821 | " --E0s=\"average\" \\\n", 822 | " --energy_key=\"energy_xtb\" \\\n", 823 | " --forces_key=\"forces_xtb\" \\\n", 824 | " --lr=0.01 \\\n", 825 | " --scaling=\"rms_forces_scaling\" \\\n", 826 | " --batch_size=10 \\\n", 827 | " --max_num_epochs=50 \\\n", 828 | " --ema \\\n", 829 | " --ema_decay=0.99 \\\n", 830 | " --amsgrad \\\n", 831 | " --default_dtype=\"float64\" \\\n", 832 | " --device=cuda \\\n", 833 | " --seed=3" 834 | ] 835 | }, 836 | { 837 | "cell_type": "code", 838 | "execution_count": null, 839 | "metadata": {}, 840 | "outputs": [], 841 | "source": [ 842 | "mace_calc = MACECalculator(model_paths=['finetuned_MACE_compiled.model'], device='cuda')\n", 843 | "\n", 844 | "init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n", 845 | "simpleMD(init_conf, temp=1200, calc=mace_calc, fname='moldyn/mace_finetuned_md.xyz', s=10, T=2000)" 846 | ] 847 | }, 848 | { 849 | "cell_type": "code", 850 | "execution_count": null, 851 | "metadata": {}, 852 | "outputs": [], 853 | "source": [ 854 | "from matplotlib import pyplot as plt\n", 855 | "from aseMolec import anaAtoms as aa\n", 856 | "\n", 857 | "tag = 'OO_intra' # 'OO_intra' #choose one of 'HH_intra', 'HC_intra', 'HO_intra', 'CC_intra', 'CO_intra', 'OO_intra'\n", 858 | "\n", 859 | "for f in ['xtb_md', 'mace_finetuned_md']:\n", 860 | " traj = read('moldyn/'+f+'.xyz', '50:') #ignore first 50 frames\n", 861 | " for at in traj:\n", 862 | " at.pbc = True #create a fake box for rdf compatibility\n", 863 | " at.cell = [100,100,100]\n", 864 | " rdf = aa.compute_rdfs_traj_avg(traj, rmax=5, nbins=70) #aseMolec provides functionality to compute RDFs\n", 865 | " plt.plot(rdf[1], rdf[0][tag], '.-', label=f, alpha=0.7, linewidth=3)\n", 866 | "\n", 867 | "plt.legend();\n", 868 | "plt.yticks([]);\n", 869 | "plt.xlabel(r'R ($\\rm \\AA$)');\n", 870 | "plt.ylabel('RDF '+tag);" 871 | ] 872 | }, 873 | { 874 | "cell_type": "code", 875 | "execution_count": null, 876 | "metadata": {}, 877 | "outputs": [], 878 | "source": [ 879 | "from ase.io import read, write\n", 880 | "import nglview as nv\n", 881 | "\n", 882 | "traj = read('moldyn/mace_finetuned_md.xyz', ':')\n", 883 | "view = nv.show_asetraj(traj)\n", 884 | "view._set_size(w='100%', h='500px')\n", 885 | "view" 886 | ] 887 | }, 888 | { 889 | "cell_type": "markdown", 890 | "metadata": {}, 891 | "source": [ 892 | "What are the results - does it work?\n", 893 | "\n", 894 | "Consider comparing more than one of the rdfs, and looking at how the MP0-model did (as we saw above) vs the finetuned version." 895 | ] 896 | }, 897 | { 898 | "cell_type": "code", 899 | "execution_count": null, 900 | "metadata": { 901 | "editable": true, 902 | "slideshow": { 903 | "slide_type": "" 904 | }, 905 | "tags": [] 906 | }, 907 | "outputs": [], 908 | "source": [] 909 | }, 910 | { 911 | "cell_type": "code", 912 | "execution_count": null, 913 | "metadata": {}, 914 | "outputs": [], 915 | "source": [] 916 | }, 917 | { 918 | "cell_type": "code", 919 | "execution_count": null, 920 | "metadata": {}, 921 | "outputs": [], 922 | "source": [] 923 | } 924 | ], 925 | "metadata": { 926 | "deepnote": { 927 | "is_reactive": false 928 | }, 929 | "deepnote_execution_queue": [], 930 | "deepnote_notebook_id": "2a16e1e8a7ad42a8825007f5cff15d19", 931 | "kernelspec": { 932 | "display_name": "Python 3 (ipykernel)", 933 | "language": "python", 934 | "name": "python3" 935 | }, 936 | "language_info": { 937 | "codemirror_mode": { 938 | "name": "ipython", 939 | "version": 3 940 | }, 941 | "file_extension": ".py", 942 | "mimetype": "text/x-python", 943 | "name": "python", 944 | "nbconvert_exporter": "python", 945 | "pygments_lexer": "ipython3", 946 | "version": "3.10.14" 947 | } 948 | }, 949 | "nbformat": 4, 950 | "nbformat_minor": 4 951 | } 952 | -------------------------------------------------------------------------------- /T03-MACE-Theory.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "!pip install mace-torch" 10 | ] 11 | }, 12 | { 13 | "cell_type": "markdown", 14 | "metadata": {}, 15 | "source": [ 16 | "# Deep Dive into the MACE Architecture" 17 | ] 18 | }, 19 | { 20 | "cell_type": "markdown", 21 | "metadata": { 22 | "editable": true, 23 | "slideshow": { 24 | "slide_type": "" 25 | }, 26 | "tags": [] 27 | }, 28 | "source": [ 29 | "In this tutorial, you will do a dive into the `MACE` code, which is a highly accurate and efficient MLIP. If you would like to understand this method in more detail, you can find the [original method paper](https://proceedings.neurips.cc/paper_files/paper/2022/file/4a36c3c51af11ed9f34615b81edb5bbc-Paper-Conference.pdf). MACE is a Message Passing Neural Network (MPNNs) Interatomic Potential that forms equivariant many body messages.\n", 30 | "\n", 31 | "MACE was developed by unifying the Atomic Cluster Expansion (ACE) approach with the equivariant MPNNs. The mathematical formalism which unifies these methods is explained in the [accompaning paper](https://doi.org/10.48550/arXiv.2205.06643). Another [useful reference](\n", 32 | "https://doi.org/10.1063/5.0155322) showcases the methods performance on published benchmark datasets aswell as updated set of equations that we will follow in this notebook. The [code implementation](https://github.com/ACEsuit/mace) is publically available and [here](https://mace-docs.readthedocs.io/en/latest/) you can find the accompaning documentation.\n", 33 | "\n", 34 | "The tutorial was developed by Will Baldwin and is based on a 'developer tutorial' from [Ilyes Batatia at University of Cambridge](https://github.com/ilyes319/mace-tutorials/)\n", 35 | "\n", 36 | "There isn't much 'computing' to be done for this notebook - the idea is to expose some of the inner functions and give you the option ot see whats going on by fiddling with the internals." 37 | ] 38 | }, 39 | { 40 | "cell_type": "markdown", 41 | "metadata": {}, 42 | "source": [ 43 | "![alt text](figures/sketch.png)" 44 | ] 45 | }, 46 | { 47 | "cell_type": "markdown", 48 | "metadata": {}, 49 | "source": [ 50 | "MACE is a function which takes in an atomic environment and computes an energy. In some of the lectures, we discussed the steps of 'designing features' and 'fitting' seperately. In MACE, these two steps are really blended together. \n", 51 | "\n", 52 | "This is illustrated in the figure above - most of the weights (and most of the computational effort) in a MACE model is constructing the atomic features. One you have these, the energy is a relatively simple learnable function of these features. \n", 53 | "\n", 54 | "In this tutoral we will look inside the feature construction of MACE and examine the individual equations and blocks of code. " 55 | ] 56 | }, 57 | { 58 | "cell_type": "markdown", 59 | "metadata": {}, 60 | "source": [ 61 | "## Schematic of MACE" 62 | ] 63 | }, 64 | { 65 | "cell_type": "markdown", 66 | "metadata": {}, 67 | "source": [ 68 | "![alt text](figures/mace_architecture.png)" 69 | ] 70 | }, 71 | { 72 | "cell_type": "markdown", 73 | "metadata": {}, 74 | "source": [ 75 | "The Figure shows a schematic of MACE. The key steps are:\n", 76 | " - 1) The Embedding. Atomic structures are turned into an initial set of 'features' describing the chemical species of each node, and describing the length and orientation of each edge. \n", 77 | " - 2) Feature Construction\n", 78 | " - 2.1) Interaction: Information is pooled from neighbours\n", 79 | " - 2.2) Product: the information aggregated from neighbours is raised to a power, creating many body descriptors.\n", 80 | " - 2.3) Update: the previous node features are updated based on the output from the product step.\n", 81 | " - 3) Readout. The new features are mapped to an energy. \n", 82 | " - 4) Repeat. The process is repeated, but the informatino describing each node (node features 1) is now much richer than it was in the previous iteration (node features 0)." 83 | ] 84 | }, 85 | { 86 | "cell_type": "markdown", 87 | "metadata": {}, 88 | "source": [ 89 | "## Default Model Parameters\n", 90 | "\n", 91 | "declare some things now to keep the notebook tidy - we will come back to them later" 92 | ] 93 | }, 94 | { 95 | "cell_type": "code", 96 | "execution_count": null, 97 | "metadata": { 98 | "id": "EqGr9Qz-lWaB" 99 | }, 100 | "outputs": [], 101 | "source": [ 102 | "import numpy as np\n", 103 | "import torch\n", 104 | "import torch.nn.functional\n", 105 | "from e3nn import o3\n", 106 | "from matplotlib import pyplot as plt\n", 107 | "import ase.io\n", 108 | "%matplotlib inline\n", 109 | "from ase.visualize import view\n", 110 | "from scipy.spatial.transform import Rotation\n", 111 | "\n", 112 | "from mace import data, modules, tools\n", 113 | "from mace.tools import torch_geometric\n", 114 | "torch.set_default_dtype(torch.float64)" 115 | ] 116 | }, 117 | { 118 | "cell_type": "code", 119 | "execution_count": null, 120 | "metadata": { 121 | "colab": { 122 | "base_uri": "https://localhost:8080/" 123 | }, 124 | "id": "42-l41XanAv2", 125 | "outputId": "4b52ee17-8acd-4eef-e22b-1c4a0776a064" 126 | }, 127 | "outputs": [], 128 | "source": [ 129 | "# setup some default prameters\n", 130 | "z_table = tools.AtomicNumberTable([1, 6, 8])\n", 131 | "atomic_energies = np.array([-1.0, -3.0, -5.0], dtype=float)\n", 132 | "cutoff = 3\n", 133 | "\n", 134 | "default_model_config = dict(\n", 135 | " num_elements=3, # number of chemical elements\n", 136 | " atomic_energies=atomic_energies, # atomic energies used for normalisation\n", 137 | " avg_num_neighbors=8, # avg number of neighbours of the atoms, used for internal normalisation of messages\n", 138 | " atomic_numbers=z_table.zs, # atomic numbers, used to specify chemical element embeddings of the model\n", 139 | " r_max=cutoff, # cutoff\n", 140 | " num_bessel=8, # number of radial features\n", 141 | " num_polynomial_cutoff=6, # smoothness of the radial cutoff\n", 142 | " max_ell=2, # expansion order of spherical harmonic adge attributes\n", 143 | " num_interactions=2, # number of layers, typically 2\n", 144 | " interaction_cls_first=modules.interaction_classes[\n", 145 | " \"RealAgnosticResidualInteractionBlock\"\n", 146 | " ], # interation block of first layer\n", 147 | " interaction_cls=modules.interaction_classes[\n", 148 | " \"RealAgnosticResidualInteractionBlock\"\n", 149 | " ], # interaction block of subsequent layers\n", 150 | " hidden_irreps=o3.Irreps(\"8x0e + 8x1o\"), # 8: number of embedding channels, 0e, 1o is specifying which equivariant messages to use. Here up to L_max=1\n", 151 | " correlation=3, # correlation order of the messages (body order - 1)\n", 152 | " MLP_irreps=o3.Irreps(\"16x0e\"), # number of hidden dimensions of last layer readout MLP\n", 153 | " gate=torch.nn.functional.silu, # nonlinearity used in last layer readout MLP\n", 154 | " )\n", 155 | "default_model = modules.MACE(**default_model_config)" 156 | ] 157 | }, 158 | { 159 | "cell_type": "markdown", 160 | "metadata": {}, 161 | "source": [ 162 | "# Representing Spherical Tensors: Ylm's\n", 163 | "\n", 164 | "Within many popular MLIP architectures, including MACE, descriptors are represented as 'spherical tensors'.\n", 165 | "\n", 166 | "To explain spherical tensors we first need to state that the spherical harmoincs 'transform' in a special way when their argument is rotated. This means that if you take a vector $\\boldsymbol{\\hat{r}}$ and rotate it by some matrix, $\\boldsymbol{\\hat{r}} \\rightarrow R\\boldsymbol{\\hat{r}}$, the value of the spherical harmonic $Y^{m}_{l} (\\boldsymbol{\\hat{r}})$ will also change. You don't need to know the maths of how the $Y_{m}^l$ transforms today, but in short this is the result: \n", 167 | "\n", 168 | "$$ Y_{lm}(R\\mathbf{r}) = \\sum_{m'} D(R)^l_{mm'} Y_{lm'}(\\mathbf{r})$$\n", 169 | "\n", 170 | "In which $D(R)$ is a unitary matrix called a Wiger D-matrix. You can find more info [here](https://en.wikipedia.org/wiki/Spherical_harmonics#Rotations).\n", 171 | "\n", 172 | "The embeddings in MACE contain a function to compute the spherical harmonics. Lets quickly demonstrate how $Y^{m}_{l}$ changes as $\\boldsymbol{\\hat{r}}$ is rotated." 173 | ] 174 | }, 175 | { 176 | "cell_type": "code", 177 | "execution_count": null, 178 | "metadata": {}, 179 | "outputs": [], 180 | "source": [ 181 | "# a function for Ylms where we evaluate for l=0,1,2.\n", 182 | "spherical_harmonics = o3.SphericalHarmonics([0,1,2], True)\n", 183 | "\n", 184 | "# evaulate spherical harmonics on a vector\n", 185 | "vector = torch.tensor([1.0, 0.2, 0.75])\n", 186 | "print(spherical_harmonics(vector))" 187 | ] 188 | }, 189 | { 190 | "cell_type": "markdown", 191 | "metadata": {}, 192 | "source": [ 193 | "Why is the array 9 elements long? We calculated the $l=0,1,2$ for all valid $m's$, and stored them in this format: $$[Y_0^0, \\ \\ Y_{1}^{-1}, Y_{1}^{0},Y_{1}^{1}, \\ \\ Y_{2}^{-2}, Y_{2}^{-1}, Y_{2}^{0}, Y_{2}^{1}, Y_{2}^{2}, ]$$\n", 194 | "\n", 195 | "Which has 9 elements in total. Generally, this is how most arrays are stored in MACE. The final dimesion of the array combines the $l$ and $m$ indicices, and concatenates the values like this. \n", 196 | "\n", 197 | "How do these transform under rotations? We can simply rotate the vector and get an idea for what happens by plotting the result:" 198 | ] 199 | }, 200 | { 201 | "cell_type": "code", 202 | "execution_count": null, 203 | "metadata": {}, 204 | "outputs": [], 205 | "source": [ 206 | "# make a list of rotated versions of the vector\n", 207 | "rotated_vectors = []\n", 208 | "vector = np.array([1.0, 0.2, 0.75])\n", 209 | "N=360\n", 210 | "for i in range(N):\n", 211 | " rotation_matrix = Rotation.from_rotvec(i * 2*np.pi * np.array([0, 0.7071, 0.7071])/360).as_matrix() # rotate around the vector [0, 0.7071, 0.7071]\n", 212 | " rotated_vectors.append(rotation_matrix @ vector)\n", 213 | "\n", 214 | "# convert to torch tensor\n", 215 | "rotated_vectors = torch.tensor(rotated_vectors, dtype=torch.float32)\n", 216 | "\n", 217 | "# compute the spherical harmonics for each vector\n", 218 | "spherical_harmonic_values = spherical_harmonics(rotated_vectors)\n", 219 | "print('shape of Y_lms array is ', spherical_harmonic_values.shape)\n", 220 | "\n", 221 | "# plot\n", 222 | "labels = [[f'l={l}, m={m}' for m in range(-l,l+1)] for l in range(3)]\n", 223 | "labels = [x for xs in labels for x in xs] # flatten\n", 224 | "plt.plot(spherical_harmonic_values.numpy(), label=labels)\n", 225 | "plt.legend()\n", 226 | "plt.xlabel('rotation angle')\n", 227 | "plt.ylabel('spherical harmonic value')\n", 228 | "plt.show()" 229 | ] 230 | }, 231 | { 232 | "cell_type": "markdown", 233 | "metadata": {}, 234 | "source": [ 235 | "The key points to note are that the $l=0$ component is constant - it does not change when the vector is rotated. The other components do change, and as $l$ increases the frequency of the features increases. " 236 | ] 237 | }, 238 | { 239 | "cell_type": "markdown", 240 | "metadata": {}, 241 | "source": [ 242 | "## Getting Invariant Quantities with `e3nn`\n", 243 | "\n", 244 | "Enevtually we want to make a model where the output is invariant to rotations. To do this, we could just take the $l=0$ piece of the spherical harmonics, but doing this doesn't give us much information. \n", 245 | "\n", 246 | "For example: Lets take a pair of vectors, get the spherical harmonics for both." 247 | ] 248 | }, 249 | { 250 | "cell_type": "code", 251 | "execution_count": null, 252 | "metadata": {}, 253 | "outputs": [], 254 | "source": [ 255 | "np.random.seed(0)\n", 256 | "vector1 = np.random.randn(3)\n", 257 | "vector1 = vector1 / np.linalg.norm(vector1)\n", 258 | "vector2 = np.random.randn(3)\n", 259 | "vector2 = vector2 / np.linalg.norm(vector2)\n", 260 | "\n", 261 | "spherical_harmonics_1 = spherical_harmonics(torch.tensor(vector1))\n", 262 | "spherical_harmonics_2 = spherical_harmonics(torch.tensor(vector2))\n", 263 | "\n", 264 | "print('l=0 component for vector 1:', spherical_harmonics_1[0])\n", 265 | "print('l=0 component for vector 2:', spherical_harmonics_2[0])" 266 | ] 267 | }, 268 | { 269 | "cell_type": "markdown", 270 | "metadata": {}, 271 | "source": [ 272 | "This isn't very useful, because the $l=0$ piece doesn't contain any angular information for either vector, so we have no way to know, for instance, the angle between them. \n", 273 | "\n", 274 | "You can test this by changing one or both of the vectors (via the numpy seed). The $l=0$ piece won't change. " 275 | ] 276 | }, 277 | { 278 | "cell_type": "markdown", 279 | "metadata": {}, 280 | "source": [ 281 | "### How do we get Invariants which describe angular information?\n", 282 | "\n", 283 | "To get a more descriptive invariant quantity, we need to do some operations on the spherical harmoincs. We care about how spherical harmonics change when you rotate the input, because its easy to keep track of how rotation affects things. This means that its easy to get back to an invariant quantity when we need to. In MACE, and many other MLIPs, this maths is done by a package called `e3nn`.\n", 284 | "\n", 285 | "`e3nn` provides functions which perform operations on spherical tensors (things with elements like $[Y_{0}^0, Y_{1}^{-1}, Y_{1}^0, ...]$), while keeping track of the rotations. One example operation is a tensor product, which takes two arrays, $A_{lm}$ and $B_{lm}$, and multiplies them to give $C_{lm}$:\n", 286 | "\n", 287 | "$$[A_{lm}] \\ \\otimes \\ [B_{lm}] \\ = \\ [C_{lm}]$$\n", 288 | "\n", 289 | "The key is that $C$ is still indexed by $l$ and $m$, so if we look at the $l=0$ piece, it will still be invariant! This means we can do a load of operations to combine spherical harmonics, and then create invariant descrpitors which know about things like angles ebtween vectors. \n", 290 | "\n", 291 | "We can demonstrate this by the two vectors above, doing the tensor product of them, and keeping all the outputs which are invariant to rotations. " 292 | ] 293 | }, 294 | { 295 | "cell_type": "code", 296 | "execution_count": null, 297 | "metadata": {}, 298 | "outputs": [], 299 | "source": [ 300 | "# set up a tensor product. \n", 301 | "# This does the product of two l=0,1,2 arrays, and maps the result to three l=0 values. \n", 302 | "tensor_product = o3.FullyConnectedTensorProduct(\n", 303 | " irreps_in1=o3.Irreps(\"1x0e + 1x1o + 1x2e\"),\n", 304 | " irreps_in2=o3.Irreps(\"1x0e + 1x1o + 1x2e\"), \n", 305 | " irreps_out=o3.Irreps(\"3x0e\"),\n", 306 | " internal_weights=False\n", 307 | ") \n", 308 | "print(tensor_product)\n", 309 | "tensor_product.visualize()" 310 | ] 311 | }, 312 | { 313 | "cell_type": "markdown", 314 | "metadata": {}, 315 | "source": [ 316 | "The diagram shows schematically how the tensor product is combining the elements of the inputs to make the output." 317 | ] 318 | }, 319 | { 320 | "cell_type": "code", 321 | "execution_count": null, 322 | "metadata": {}, 323 | "outputs": [], 324 | "source": [ 325 | "# product the arrays\n", 326 | "product = tensor_product(\n", 327 | " spherical_harmonics_1.unsqueeze(0), \n", 328 | " spherical_harmonics_2.unsqueeze(0), \n", 329 | " weight=torch.arange(1,10,1) # the product has weights which can be trained - for now I have fixed them\n", 330 | ")\n", 331 | "print('invariant outputs:', product)" 332 | ] 333 | }, 334 | { 335 | "cell_type": "markdown", 336 | "metadata": {}, 337 | "source": [ 338 | "Now we have three outputs, and they have the property that they are invariant to rotations of the *entire* structure (i.e. rotating oth vectors) but then still change when the angle between the vectors changes. \n", 339 | "\n", 340 | "We can demonstrate this:" 341 | ] 342 | }, 343 | { 344 | "cell_type": "code", 345 | "execution_count": null, 346 | "metadata": {}, 347 | "outputs": [], 348 | "source": [ 349 | "angle = 77.7 # degrees\n", 350 | "rotation_matrix = Rotation.from_rotvec(angle * 2*np.pi * np.array([0, 0.7071, 0.7071])/360).as_matrix()\n", 351 | "\n", 352 | "rotated_vec1 = rotation_matrix @ vector1\n", 353 | "rotated_vec2 = rotation_matrix @ vector2\n", 354 | "\n", 355 | "# get the spherical harmpnics\n", 356 | "spherical_harmonics_1 = spherical_harmonics(torch.from_numpy(rotated_vec1))\n", 357 | "spherical_harmonics_2 = spherical_harmonics(torch.from_numpy(rotated_vec2))\n", 358 | "\n", 359 | "product = tensor_product(\n", 360 | " spherical_harmonics_1.unsqueeze(0), \n", 361 | " spherical_harmonics_2.unsqueeze(0), \n", 362 | " weight=torch.arange(1,10,1) # the product has weights which can be trained - for now I have fixed them\n", 363 | ")\n", 364 | "print('invariant outputs:', product)" 365 | ] 366 | }, 367 | { 368 | "cell_type": "markdown", 369 | "metadata": {}, 370 | "source": [ 371 | "You can see that the values are the same as before the rotation, but if you go back and change the seed used for creating the two vectors, they will be different. This meas that we have have made invariant quantities which can tell two pairs of unit vectors apart. \n", 372 | "\n", 373 | "These kind of operations are what MACE is built on." 374 | ] 375 | }, 376 | { 377 | "cell_type": "markdown", 378 | "metadata": {}, 379 | "source": [ 380 | "# MACE Feature Construction\n", 381 | "\n", 382 | "We can now go through some key parts of the feature construction and look at the internal variables. The schematic at the top of the document has a full overview, and the key bits will be repeated. \n", 383 | "\n", 384 | "## 0. Data Prep\n", 385 | "\n", 386 | "The first step in MACE is to take the atomic structure and represent it as a list of atoms and 'edges'. An edge is simply a connection between two atoms. The cutoff radius of the MACE layer determines which atoms are connected by an edge. \n", 387 | "\n", 388 | "Below we load a list of molecule, and show how the atoms and edges are stored." 389 | ] 390 | }, 391 | { 392 | "cell_type": "code", 393 | "execution_count": null, 394 | "metadata": {}, 395 | "outputs": [], 396 | "source": [ 397 | "single_molecule = ase.io.read('data/solvent_rotated.xyz', index='0')\n", 398 | "\n", 399 | "Rcut = 3.0 # cutoff radius\n", 400 | "z_table = tools.AtomicNumberTable([1, 6, 8])\n", 401 | "\n", 402 | "config = data.Configuration(\n", 403 | " atomic_numbers=single_molecule.numbers,\n", 404 | " positions=single_molecule.positions\n", 405 | ")\n", 406 | "\n", 407 | "# we handle configurations using the AtomicData class\n", 408 | "batch = data.AtomicData.from_config(config, z_table=z_table, cutoff=Rcut)\n", 409 | "print(\"positions:\", batch.positions)\n", 410 | "print(\"node_attrs:\", batch.node_attrs)\n", 411 | "print(\"edge_index:\", batch.edge_index)" 412 | ] 413 | }, 414 | { 415 | "cell_type": "markdown", 416 | "metadata": {}, 417 | "source": [ 418 | "The `node_attrs` describe the species. `[1,0,0]` means the first kind of element, `[0,1,0]` means the second kind, and so on. \n", 419 | "\n", 420 | "In this case we told MACE to work with Hydrogen, Carbon and Oxygen with the line `z_table = ...`.\n", 421 | "\n", 422 | "The 'edges' are stored by a list of 'senders' and 'receievers', which represent the start and end point of each edge. \n", 423 | "\n", 424 | "We can now compute the lengths and directions of each edge:" 425 | ] 426 | }, 427 | { 428 | "cell_type": "code", 429 | "execution_count": null, 430 | "metadata": {}, 431 | "outputs": [], 432 | "source": [ 433 | "vectors, lengths = modules.utils.get_edge_vectors_and_lengths(\n", 434 | " positions=batch[\"positions\"],\n", 435 | " edge_index=batch[\"edge_index\"],\n", 436 | " shifts=batch[\"shifts\"],\n", 437 | ")\n", 438 | "print(f'there are {batch.positions.shape[0]} nodes and {len(lengths)} edges')\n", 439 | "print(f'lengths is shape {lengths.shape}')\n", 440 | "print(f'vectors is shape {vectors.shape}')" 441 | ] 442 | }, 443 | { 444 | "cell_type": "markdown", 445 | "metadata": {}, 446 | "source": [ 447 | "\n", 448 | "## 1. Embeddings\n", 449 | "\n", 450 | "We now take this information and make the initial inputs to the MACE model. This is called the embedding. \n", 451 | "\n", 452 | "#### Key concept: MACE 'channels'\n", 453 | "\n", 454 | "The equation for the initial node features looks like this: \n", 455 | "$$h_{i,k00}^{(0)} = \\sum_z W_{kz} \\delta_{zz_{i}}$$\n", 456 | "Take the exmaple where all atoms are either hydrogen or carbon. If atom $i$ is hydrogen, the its initial node features are just $W_{k0}$. If atom $n$ is carbon it will get initial features $W_{k1}$. This means that each atom is given a vector of length $K$, based on its chemical species. \n", 457 | "\n", 458 | "```In mace, the range of k is referred to as the number of 'channels'. This is the fundamental 'size' of the descriptor```\n", 459 | "the figure at the top of the notebook sort of shows this. \n", 460 | "\n", 461 | "This operation of giving each atom a length-K vector depending on its species is called the node 'embedding'. \n", 462 | "\n", 463 | "We also 'embed' the lengths and directions of the edges. The lengths of the edges are mapped through a set of 8 bessel functions, and for the dircetion of the edge we just calculate the spherical harmoincs.\n", 464 | "\n", 465 | "![alt text](figures/embedding_layer.png)" 466 | ] 467 | }, 468 | { 469 | "cell_type": "code", 470 | "execution_count": null, 471 | "metadata": {}, 472 | "outputs": [], 473 | "source": [ 474 | "# set up a mace model to get all of the blocks in one place:\n", 475 | "model = modules.MACE(**default_model_config)\n", 476 | "\n", 477 | "initial_node_features = model.node_embedding(batch.node_attrs)\n", 478 | "edge_features = model.radial_embedding(lengths, batch[\"node_attrs\"], batch[\"edge_index\"], z_table)\n", 479 | "edge_attributes = model.spherical_harmonics(vectors)\n", 480 | "\n", 481 | "print('initial_node_features is (num_atoms, num_channels):', initial_node_features.shape)\n", 482 | "print('edge_features is (num_edge, num_bessel_func):', edge_features.shape)\n", 483 | "print('edge_attributes is (num_edge, dimension of spherical harmonics):', edge_attributes.shape)\n", 484 | "print(\n", 485 | " '\\nInitial node features. Note that they are the same for each chemical element\\n', \n", 486 | " initial_node_features\n", 487 | ")" 488 | ] 489 | }, 490 | { 491 | "cell_type": "code", 492 | "execution_count": null, 493 | "metadata": {}, 494 | "outputs": [], 495 | "source": [ 496 | "# we can plot what the 'edge features' look like as a function of edge length: \n", 497 | "dists = torch.tensor(np.linspace(0.1, cutoff, 100), dtype=torch.get_default_dtype()).unsqueeze(-1)\n", 498 | "radials = model.radial_embedding(dists, batch[\"node_attrs\"], batch[\"edge_index\"], z_table)\n", 499 | "\n", 500 | "for i in range(radials.shape[1]):\n", 501 | " plt.plot(dists, radials[:, i], label=f'edge feat {i}')\n", 502 | "\n", 503 | "# Add title, labels, and legend\n", 504 | "plt.title(\"Edge features\")\n", 505 | "plt.xlabel(\"distance / A\")\n", 506 | "plt.ylabel(\"Value\")\n", 507 | "plt.legend()\n", 508 | "\n", 509 | "# Display the plot\n", 510 | "plt.show()" 511 | ] 512 | }, 513 | { 514 | "cell_type": "markdown", 515 | "metadata": {}, 516 | "source": [ 517 | "The parameter `max_ell` in the mace config (top of file) controls the maximum $l$ in the spherical harmonics, and hence the dimension of the `edge_attributes`.\n", 518 | "\n", 519 | "The intial node features are invariant, and depend only on the speices. You can check this by printing them and comparing to `batch.node_attrs`.\n", 520 | "\n", 521 | "The edge features are invariant because they only depend on length, and the edge attributes are of course spherical harmonics, and change as things rotate." 522 | ] 523 | }, 524 | { 525 | "cell_type": "markdown", 526 | "metadata": {}, 527 | "source": [ 528 | "\n", 529 | "## 2. Interaction\n", 530 | "\n", 531 | "Having created the intial features, the fun begins. The interaction block is tasked with pooling information over neighbours while tracking the properties of all arrays with $(l,m)$ indices. The output is a set of atomic features which contain information about the neighbours but are still 2-body. \n", 532 | "\n", 533 | "The figure below shows all the equations of the interaction which is used in the first layer. In the second layer of mace, the interaction step is significantly more complex - the end of this notebook discusses this if you are interested. \n", 534 | "\n", 535 | "The first layer of MACE is actually very similar to the simpler, and related, ACE model. We won't go through the internals of the interaction, but its at the end of the notebook if you are interested. \n", 536 | "\n", 537 | "![alt text](figures/interaction_first_layer.png)" 538 | ] 539 | }, 540 | { 541 | "cell_type": "code", 542 | "execution_count": null, 543 | "metadata": {}, 544 | "outputs": [], 545 | "source": [ 546 | "Interaction = model.interactions[0]\n", 547 | "\n", 548 | "intermediate_node_features, sc = Interaction(\n", 549 | " node_feats=initial_node_features,\n", 550 | " node_attrs=batch.node_attrs,\n", 551 | " edge_feats=edge_features,\n", 552 | " edge_attrs=edge_attributes,\n", 553 | " edge_index=batch.edge_index,\n", 554 | ")\n", 555 | "\n", 556 | "print(\"the output of the interaction is (num_atoms, channels, dim. spherical harmonics):\", intermediate_node_features.shape)" 557 | ] 558 | }, 559 | { 560 | "cell_type": "markdown", 561 | "metadata": {}, 562 | "source": [ 563 | "One important part of the interaction is how the `edge_attrs` (the bessel function values) enter. \n", 564 | "\n", 565 | "The bessel function values are mapped through a multilayer perceptron ('radial MLP') to make a large number of learnable functions of edge length. The output is sometimes called a 'learnable raidal basis' $R_{kl}(r_{ij})$ - one radial function for each $(k,l)$ combination.\n", 566 | "\n", 567 | "These learnable functions are then combined with the spherical harmonics of the $j\\rightarrow i$ edge, and the initial node features of the neighbouring atom ($j$). The result is that $\\phi_{ij,klm}(\\mathbf{r}_{ij})$ is quite a rich set of descriptors of the $j$'th neighbour atom. \n", 568 | "\n", 569 | "We can plot some of these radial functions to see how kind of shapes they have." 570 | ] 571 | }, 572 | { 573 | "cell_type": "code", 574 | "execution_count": null, 575 | "metadata": {}, 576 | "outputs": [], 577 | "source": [ 578 | "# visulaise what the radial MLP does\n", 579 | "dists = torch.tensor(np.linspace(0.1, cutoff, 100), dtype=torch.get_default_dtype()).unsqueeze(-1)\n", 580 | "\n", 581 | "# first do radial embedding again\n", 582 | "edge_feats_scan = model.radial_embedding(dists, batch[\"node_attrs\"], batch[\"edge_index\"], model.atomic_numbers)\n", 583 | "\n", 584 | "# put the edge feats into the radial MLP\n", 585 | "tp_weights_scan = model.interactions[0].conv_tp_weights(edge_feats_scan).detach().numpy()\n", 586 | "print('the output of the radial MLP is (num_dists, num_weights):', tp_weights_scan.shape)\n", 587 | "\n", 588 | "# plot the outputs\n", 589 | "num_basis_to_print = 5\n", 590 | "for i in range(num_basis_to_print):\n", 591 | " plt.plot(dists, tp_weights_scan[:, i], label=f'Learnable Radial {i}')\n", 592 | "\n", 593 | "# Add title, labels, and legend\n", 594 | "plt.title(f\"First {num_basis_to_print} MACE learnable radial functions (untrained)\")\n", 595 | "plt.xlabel(\"distance / A\")\n", 596 | "plt.ylabel(\"Value\")\n", 597 | "plt.legend()\n", 598 | "\n", 599 | "# Display the plot\n", 600 | "plt.show()" 601 | ] 602 | }, 603 | { 604 | "cell_type": "markdown", 605 | "metadata": {}, 606 | "source": [ 607 | "Of course these functional forms will change after training. " 608 | ] 609 | }, 610 | { 611 | "cell_type": "markdown", 612 | "metadata": {}, 613 | "source": [ 614 | "\n", 615 | "## 3. Product\n", 616 | "\n", 617 | "The key operation of MACE is the efficient construction of higher order features from the ${A}_{i}^{(t)}$-features (the output of the interaction).\n", 618 | "This is achieved by first forming tensor products of the features, and then symmetrising, as shown in the figure. This operation:\n", 619 | "$$\n", 620 | "{B}^{(t)}_{i,\\eta_{\\nu} k LM}\n", 621 | "= \\sum_{{l}{m}} \\mathcal{C}^{LM}_{\\eta_{\\nu}, l m} \\prod_{\\xi = 1}^{\\nu} A_{i,k l_\\xi m_\\xi}^{(t)}\n", 622 | "$$\n", 623 | "is essentially doing $A_{lm} \\otimes A_{lm} \\otimes ...$ and eventually getting back to something which transforms in a known way - ${B}^{(t)}_{i,\\eta_{\\nu} k LM}$. It happens on each node - all pooling between neighbours is done in the interction.\n", 624 | "\n", 625 | "Like in the example above with two vectors, this product operation creates features which have angular information, but which are still invariant. \n", 626 | "\n", 627 | "![alt text](figures/product_layer.png)" 628 | ] 629 | }, 630 | { 631 | "cell_type": "markdown", 632 | "metadata": {}, 633 | "source": [ 634 | "The final part of the product is to linearly mix everything together and create a new set of node features.\n", 635 | "\n", 636 | "These node features don't have to be invariant. Its possible to retain some level of equivariance and then use this extra information in the second layer of MACE, where the whole interaction an product is repeated. \n", 637 | "\n", 638 | "Whether to do this is controlled by the `max_L` parameter to MACE (see the model config). If this is set to 0, then only invariant features are retained. If it is set to 1 (which is the case here), the new node features will have $l=0$ and $l=1$ pieces. " 639 | ] 640 | }, 641 | { 642 | "cell_type": "code", 643 | "execution_count": null, 644 | "metadata": {}, 645 | "outputs": [], 646 | "source": [ 647 | "new_node_features = model.products[0](\n", 648 | " node_feats=intermediate_node_features,\n", 649 | " node_attrs=batch.node_attrs,\n", 650 | " sc=sc,\n", 651 | ")\n", 652 | "print('new node feats are (num_atoms, (L+1)^2 * num_channels):', new_node_features.shape)" 653 | ] 654 | }, 655 | { 656 | "cell_type": "markdown", 657 | "metadata": {}, 658 | "source": [ 659 | "you should find that the array is (num_atoms, 32), since 32=8*4, where 4 is because we have $[Y_{00}, Y_{1,-1}, Y_{1,0}, Y_{1,1}]$ pieces, and 8 is because we have 8 'channels'.\n", 660 | "\n", 661 | "You can check the equivariance of these features by modifying the above code to read a different config from `rotated_solvent.xyz`. This will be the same structure, but rotated. You should see that the first 32 elements are the same (since they are invariant) and the rest change." 662 | ] 663 | }, 664 | { 665 | "cell_type": "markdown", 666 | "metadata": {}, 667 | "source": [ 668 | "\n", 669 | "## 4. Readout\n", 670 | "\n", 671 | "Finally, we can take the new node features and create an actual energy. This is done by passing the node features through a `readout`. \n", 672 | "\n", 673 | "![alt text](figures/readout.png)" 674 | ] 675 | }, 676 | { 677 | "cell_type": "markdown", 678 | "metadata": {}, 679 | "source": [ 680 | "In an $S$-layer MACE, the readout from the features at layer $s$ is:\n", 681 | "\n", 682 | "\\begin{equation*}\n", 683 | " \\mathcal{R}^{(s)} \\left( \\boldsymbol{h}_i^{(s)} \\right) =\n", 684 | " \\begin{cases}\n", 685 | " \\sum_{k}W^{(s)}_{k}h^{(s)}_{i,k00} & \\text{if} \\;\\; 1 < s < S \\\\[13pt]\n", 686 | " {\\rm MLP} \\left( \\left\\{ h^{(s)}_{i,k00} \\right\\}_k \\right) &\\text{if} \\;\\; s = S\n", 687 | " \\end{cases}\n", 688 | "\\end{equation*}\n", 689 | "\n", 690 | "In our example case this maps the 32 dimensional $h^{(1)}_{i,k00}$, the invariant part os the node features after the first interaction to the first term in the aotmic site energy:" 691 | ] 692 | }, 693 | { 694 | "cell_type": "code", 695 | "execution_count": null, 696 | "metadata": {}, 697 | "outputs": [], 698 | "source": [ 699 | "print('first readout =', model.readouts[0], '\\n')\n", 700 | "\n", 701 | "energy_layer_1 = model.readouts[0](new_node_features)\n", 702 | "print('energy_layer_1:', energy_layer_1.shape)\n", 703 | "print(energy_layer_1)" 704 | ] 705 | }, 706 | { 707 | "cell_type": "markdown", 708 | "metadata": {}, 709 | "source": [ 710 | "And we have made an energy for each node!" 711 | ] 712 | }, 713 | { 714 | "cell_type": "markdown", 715 | "metadata": {}, 716 | "source": [ 717 | "## 5. Repeat\n", 718 | "\n", 719 | "The Interaction, product and readout are repeated twice, and all the atomic energy contributions are summed up to get the total energy." 720 | ] 721 | }, 722 | { 723 | "cell_type": "markdown", 724 | "metadata": {}, 725 | "source": [ 726 | "# Interaction Block in more detail\n", 727 | "\n", 728 | "At the second layer, the interaction block has a much harder task. This is because the features at the end of the first MACE layer may not be scalar, but may have an $(l,m)$ index pair. \n", 729 | "\n", 730 | "In general, the equations for layer $s$ are:\n", 731 | "\n", 732 | "![alt text](figures/interaction_layer_second.png)" 733 | ] 734 | }, 735 | { 736 | "cell_type": "markdown", 737 | "metadata": {}, 738 | "source": [ 739 | "In this case, the learnable radial functions have a more complicated funciton. \n", 740 | "\n", 741 | "These learnable functions are then used as weights in the operation that follows, which does a product between the $j\\rightarrow i$ edge and the initial features of node $j$. You can see the detail in the equations above if you are interested.\n", 742 | "\n", 743 | "Because the learnable functions are used as weights in the following `conv_tp` operation, the output shape is just the number of weights required by that block, and they are all invariants.\n", 744 | "\n", 745 | "### Task\n", 746 | "\n", 747 | "If you really want to understand whats going on, try to work out what the output shape of the radial MLP in the second layer should be. You can access the block using some of the code snippets below:" 748 | ] 749 | }, 750 | { 751 | "cell_type": "code", 752 | "execution_count": null, 753 | "metadata": {}, 754 | "outputs": [], 755 | "source": [ 756 | "# visulaise what the radial MLP does\n", 757 | "dists = torch.tensor(np.linspace(0.1, cutoff, 100), dtype=torch.get_default_dtype()).unsqueeze(-1)\n", 758 | "\n", 759 | "# first do radial embedding again\n", 760 | "edge_feats_scan = model.radial_embedding(dists, batch[\"node_attrs\"], batch[\"edge_index\"], model.atomic_numbers)\n", 761 | "\n", 762 | "# put the edge feats into the radial MLP\n", 763 | "### model.interactions[1] for the second layer!!!!\n", 764 | "# conv_tp_weights is the name for the radial MLP...\n", 765 | "tp_weights_scan = model.interactions[1].conv_tp_weights(edge_feats_scan).detach().numpy()\n", 766 | "print('the output of the radial MLP is (num_dists, num_weights):', tp_weights_scan.shape)\n", 767 | "\n", 768 | "\n", 769 | "# plot the outputs if you want\n", 770 | "num_basis_to_print = 5\n", 771 | "for i in range(num_basis_to_print):\n", 772 | " plt.plot(dists, tp_weights_scan[:, i], label=f'Learnable Radial {i}')\n", 773 | "\n", 774 | "# Add title, labels, and legend\n", 775 | "plt.title(f\"First {num_basis_to_print} MACE learnable radial functions (untrained)\")\n", 776 | "plt.xlabel(\"distance / A\")\n", 777 | "plt.ylabel(\"Value\")\n", 778 | "plt.legend()\n", 779 | "\n", 780 | "# Display the plot\n", 781 | "plt.show()" 782 | ] 783 | } 784 | ], 785 | "metadata": { 786 | "colab": { 787 | "provenance": [], 788 | "toc_visible": true 789 | }, 790 | "kernelspec": { 791 | "display_name": "Python 3 (ipykernel)", 792 | "language": "python", 793 | "name": "python3" 794 | }, 795 | "language_info": { 796 | "codemirror_mode": { 797 | "name": "ipython", 798 | "version": 3 799 | }, 800 | "file_extension": ".py", 801 | "mimetype": "text/x-python", 802 | "name": "python", 803 | "nbconvert_exporter": "python", 804 | "pygments_lexer": "ipython3", 805 | "version": "3.11.7" 806 | } 807 | }, 808 | "nbformat": 4, 809 | "nbformat_minor": 4 810 | } 811 | -------------------------------------------------------------------------------- /T04-MLIP-Apps.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "id": "1", 7 | "metadata": {}, 8 | "outputs": [], 9 | "source": [ 10 | "!pip install janus-core " 11 | ] 12 | }, 13 | { 14 | "cell_type": "code", 15 | "execution_count": 1, 16 | "id": "43d4428d", 17 | "metadata": {}, 18 | "outputs": [], 19 | "source": [ 20 | "from data import get_data\n", 21 | "get_data(filename=\"LiFePO4_supercell.cif\",folder=\"data\")" 22 | ] 23 | }, 24 | { 25 | "cell_type": "markdown", 26 | "id": "abe87e70", 27 | "metadata": {}, 28 | "source": [ 29 | "\n", 30 | "# Nudge Elastic bands for lithium iron phosphate\n", 31 | "\n", 32 | "The foundation potentials (called universal by some) provide an easy way to use MLIPs. In this tutorial we will see how well some of these models perform against each other and against the published DFT literature.\n", 33 | "\n", 34 | "A top chart of main MLIPs is available\n", 35 | "\n", 36 | "https://matbench-discovery.materialsproject.org/\n", 37 | "\n", 38 | "References:\n", 39 | "\n", 40 | "1. https://pubs.rsc.org/en/content/articlelanding/2011/ee/c1ee01782a\n", 41 | "2. https://doi.org/10.1103/PhysRevApplied.7.034007\n", 42 | "3. https://arxiv.org/abs/2401.00096\n", 43 | "4. https://arxiv.org/abs/2402.03789\n", 44 | "\n", 45 | "As a task we will determine the activation energies of Li diffusion along the [010] and [001] directions (or Paths b and c how we call them here) in lithium iron phosphate (LiFePO_4) a cathode material for lithium ion batteries. \n", 46 | " \n", 47 | "DFT references energies are: barrier heights: path b = 0.27 eV and path c = 2.5 eV. (see table 1 in https://doi.org/10.1039/C5TA05062F)\n", 48 | "\n", 49 | "We will try CHGNet, M3GNET and MACE(you can try small, medium and large variants). We start with a fully working example with CHGnet. Once you worked your way through the netbook and gain some understanding of the mechanics, duplicate the notebook and change it to use your new MLIP.\n", 50 | "\n", 51 | "**HINT** to change the MLIP change the bits of code containing\n", 52 | "\n", 53 | "```python\n", 54 | "sp = SinglePoint(\n", 55 | " struct=LFPO.copy(),\n", 56 | " architecture=\"chgnet\",\n", 57 | " device='cuda',\n", 58 | ")\n", 59 | "```\n", 60 | "\n", 61 | "to \n", 62 | "\n", 63 | "```python\n", 64 | "sp = SinglePoint(\n", 65 | " struct=LFPO.copy(),\n", 66 | " architecture=\"mace_mp\",\n", 67 | " device='cuda',\n", 68 | " calc_kwargs={'model_paths':'small','default_dtype':'float64'}\n", 69 | ")\n", 70 | "```\n", 71 | "for MACE_MP small\n", 72 | "\n", 73 | "or to \n", 74 | "\n", 75 | "```python\n", 76 | "sp = SinglePoint(\n", 77 | " struct=LFPO.copy(),\n", 78 | " architecture=\"m3gnet\",\n", 79 | " device='cuda',\n", 80 | ")\n", 81 | "```\n", 82 | "for M3GNET\n", 83 | "\n", 84 | "\n", 85 | "this notebook is slightly altered and simplified from https://github.com/materialsvirtuallab/matcalc/blob/main/examples/LiFePO4-NEB.ipynb\n", 86 | "\n", 87 | "codes used\n", 88 | "- ASE: https://gitlab.com/ase/ase\n", 89 | "- pymatgen: https://github.com/materialsproject/pymatgen\n", 90 | "- janus-core: https://github.com/stfc/janus-core" 91 | ] 92 | }, 93 | { 94 | "cell_type": "code", 95 | "execution_count": null, 96 | "id": "cf403bdc", 97 | "metadata": {}, 98 | "outputs": [], 99 | "source": [ 100 | "\n", 101 | "#packages needed\n", 102 | "\n", 103 | "from ase.io import read,write\n", 104 | "from janus_core.calculations.single_point import SinglePoint\n", 105 | "from janus_core.calculations.geom_opt import optimize\n", 106 | "from ase.mep import NEB, NEBTools\n", 107 | "\n", 108 | "from ase.optimize import LBFGS,BFGS\n", 109 | "from pymatgen.io.ase import AseAtomsAdaptor\n", 110 | "from itertools import chain\n", 111 | "\n", 112 | "from pymatgen.core import PeriodicSite, Structure\n", 113 | "\n", 114 | "from ase.visualize import view\n", 115 | "\n", 116 | "device = \"cuda\"" 117 | ] 118 | }, 119 | { 120 | "cell_type": "markdown", 121 | "id": "3", 122 | "metadata": {}, 123 | "source": [ 124 | "# Prepare NEB end structures\n" 125 | ] 126 | }, 127 | { 128 | "cell_type": "markdown", 129 | "id": "4", 130 | "metadata": {}, 131 | "source": [ 132 | "## Initial structures\n", 133 | "\n", 134 | "intial structure can be downloaded from materials project, mp-19017, here we we provide the initial structure, the supercell path b and c end structures for convenience.\n" 135 | ] 136 | }, 137 | { 138 | "cell_type": "code", 139 | "execution_count": null, 140 | "id": "5", 141 | "metadata": {}, 142 | "outputs": [], 143 | "source": [ 144 | "LFPO = read(\"data/LiFePO4_supercell.cif\")\n", 145 | "view(LFPO,viewer='ngl')" 146 | ] 147 | }, 148 | { 149 | "cell_type": "markdown", 150 | "id": "6", 151 | "metadata": {}, 152 | "source": [ 153 | "## Relax supercell" 154 | ] 155 | }, 156 | { 157 | "cell_type": "code", 158 | "execution_count": null, 159 | "id": "8", 160 | "metadata": {}, 161 | "outputs": [], 162 | "source": [ 163 | "# setup a single point\n", 164 | "sp = SinglePoint(\n", 165 | " struct=LFPO.copy(),\n", 166 | " architecture=\"chgnet\",\n", 167 | " device=device)\n", 168 | "relaxed_LFPO = optimize(\n", 169 | " struct=sp.struct,\n", 170 | " fmax=0.01)\n", 171 | "view(relaxed_LFPO,viewer='ngl')\n" 172 | ] 173 | }, 174 | { 175 | "cell_type": "markdown", 176 | "id": "9", 177 | "metadata": {}, 178 | "source": [ 179 | "## Create NEB start, end structures -- b and c directions\n" 180 | ] 181 | }, 182 | { 183 | "cell_type": "code", 184 | "execution_count": null, 185 | "id": "10", 186 | "metadata": {}, 187 | "outputs": [], 188 | "source": [ 189 | "# for end b remove site 11 \n", 190 | "# for end c remove site 4\n", 191 | "# for start bc remove site 5\n", 192 | "# NEB path along b and c directions have the same starting image.\n", 193 | "\n", 194 | "LFPO_end_b = relaxed_LFPO.copy()\n", 195 | "del LFPO_end_b[11]\n", 196 | "\n", 197 | "sp_end_b = SinglePoint(\n", 198 | " struct=LFPO_end_b.copy(),\n", 199 | " architecture=\"chgnet\",\n", 200 | " device=device,\n", 201 | ")\n", 202 | "relaxed_end_b_LFPO = optimize(\n", 203 | " struct=sp_end_b.struct,\n", 204 | " fmax=0.01,\n", 205 | " filter_func=None,\n", 206 | ")\n", 207 | "view(relaxed_end_b_LFPO,viewer='ngl')\n" 208 | ] 209 | }, 210 | { 211 | "cell_type": "code", 212 | "execution_count": null, 213 | "id": "39d057e2", 214 | "metadata": {}, 215 | "outputs": [], 216 | "source": [ 217 | "LFPO_end_c = relaxed_LFPO.copy()\n", 218 | "del LFPO_end_c[4]\n", 219 | "\n", 220 | "sp_end_c = SinglePoint(\n", 221 | " struct=LFPO_end_c.copy(),\n", 222 | " architecture=\"chgnet\",\n", 223 | " device=device,\n", 224 | ")\n", 225 | "relaxed_end_c_LFPO = optimize(\n", 226 | " struct=sp_end_c.struct,\n", 227 | " fmax=0.01,\n", 228 | " filter_func=None,\n", 229 | ")\n", 230 | "view(relaxed_end_b_LFPO,viewer=\"ngl\")\n" 231 | ] 232 | }, 233 | { 234 | "cell_type": "code", 235 | "execution_count": null, 236 | "id": "e562bd7f", 237 | "metadata": {}, 238 | "outputs": [], 239 | "source": [ 240 | "LFPO_start_bc = relaxed_LFPO.copy()\n", 241 | "del LFPO_start_bc[5]\n", 242 | "sp_start_bc = SinglePoint(\n", 243 | " struct=LFPO_start_bc.copy(),\n", 244 | " architecture=\"chgnet\",\n", 245 | " device=device)\n", 246 | "relaxed_start_bc_LFPO = optimize(\n", 247 | " struct=sp_start_bc.struct,\n", 248 | " fmax=0.01,\n", 249 | " filter_func=None,\n", 250 | ")\n", 251 | "\n", 252 | "view(relaxed_start_bc_LFPO,viewer=\"ngl\")" 253 | ] 254 | }, 255 | { 256 | "cell_type": "markdown", 257 | "id": "d721de7a-3975-4f09-9d30-9bd9ffb38c39", 258 | "metadata": {}, 259 | "source": [ 260 | "# Calculate NEB path and barriers" 261 | ] 262 | }, 263 | { 264 | "cell_type": "markdown", 265 | "id": "25fbd424-caca-44a6-84e6-ced6609b4d48", 266 | "metadata": {}, 267 | "source": [ 268 | "## Path b" 269 | ] 270 | }, 271 | { 272 | "cell_type": "code", 273 | "execution_count": null, 274 | "id": "579d17c7", 275 | "metadata": {}, 276 | "outputs": [], 277 | "source": [ 278 | "nimages = 7\n", 279 | "\n", 280 | "start = AseAtomsAdaptor.get_structure(relaxed_start_bc_LFPO)\n", 281 | "end_b = AseAtomsAdaptor.get_structure(relaxed_end_b_LFPO)\n", 282 | "images_p = start.interpolate(end_b, nimages=nimages+1, pbc=False, interpolate_lattices=False, autosort_tol=0.5)\n", 283 | "images_b = [ p.to_ase_atoms() for p in images_p]\n", 284 | "\n", 285 | "# Set calculators:\n", 286 | "for image in images_b:\n", 287 | " sp = SinglePoint(\n", 288 | " struct=image,\n", 289 | " architecture=\"chgnet\",\n", 290 | " device=device)\n", 291 | "neb_b = NEB(images_b,climb=True,allow_shared_calculator=True)\n", 292 | "\n", 293 | "# do the neb ptimize:\n", 294 | "opt = BFGS(neb_b, trajectory='neb_b.traj')\n", 295 | "opt.run(fmax=0.05)\n", 296 | "\n", 297 | "#view the final path\n", 298 | "nebtools_b = NEBTools(images_b)\n", 299 | "\n", 300 | "# Get the calculated barrier and the energy change of the reaction.\n", 301 | "Ef, dE = nebtools_b.get_barrier()\n", 302 | "\n", 303 | "# Get the barrier without any interpolation between highest images.\n", 304 | "Ef, dE = nebtools_b.get_barrier(fit=False)\n", 305 | "\n", 306 | "# Get the actual maximum force at this point in the simulation.\n", 307 | "max_force = nebtools_b.get_fmax()\n", 308 | "\n", 309 | "# Create a figure like that coming from ASE-GUI.\n", 310 | "fig_b = nebtools_b.plot_band()\n", 311 | "\n", 312 | "view(images_b, viewer=\"ngl\")" 313 | ] 314 | }, 315 | { 316 | "cell_type": "markdown", 317 | "id": "f3f9466b-2194-4e22-ba57-9fc0e13845df", 318 | "metadata": {}, 319 | "source": [ 320 | "## Path c" 321 | ] 322 | }, 323 | { 324 | "cell_type": "code", 325 | "execution_count": null, 326 | "id": "4e0a3652", 327 | "metadata": {}, 328 | "outputs": [], 329 | "source": [ 330 | "end_c = AseAtomsAdaptor.get_structure(relaxed_end_c_LFPO)\n", 331 | "images_p = start.interpolate(end_c, nimages=nimages+1, pbc=False, interpolate_lattices=False, autosort_tol=0.5)\n", 332 | "images_c = [ p.to_ase_atoms() for p in images_p]\n", 333 | "\n", 334 | "# Set calculators:\n", 335 | "for image in images_c:\n", 336 | " sp = SinglePoint(\n", 337 | " struct=image,\n", 338 | " architecture=\"chgnet\",\n", 339 | " device=device)\n", 340 | "# print(sp.run()['energy']) \n", 341 | "neb_c = NEB(images_c,climb=True,allow_shared_calculator=True)\n", 342 | "# do the nebptimize:\n", 343 | "opt = BFGS(neb_c, trajectory='neb_c.traj')\n", 344 | "opt.run(fmax=0.05)\n", 345 | "\n", 346 | "#view the final path\n", 347 | "nebtools_c = NEBTools(images_c)\n", 348 | "\n", 349 | "# Get the calculated barrier and the energy change of the reaction.\n", 350 | "Ef, dE = nebtools_c.get_barrier()\n", 351 | "\n", 352 | "# Get the barrier without any interpolation between highest images.\n", 353 | "Ef, dE = nebtools_c.get_barrier(fit=False)\n", 354 | "\n", 355 | "# Get the actual maximum force at this point in the simulation.\n", 356 | "max_force = nebtools_c.get_fmax()\n", 357 | "\n", 358 | "# Create a figure like that coming from ASE-GUI.\n", 359 | "fig_c = nebtools_c.plot_band()\n", 360 | "\n", 361 | "view(images_c,viewer=\"ngl\")" 362 | ] 363 | }, 364 | { 365 | "cell_type": "markdown", 366 | "id": "16", 367 | "metadata": {}, 368 | "source": [ 369 | "# View NEB path in one snapshot" 370 | ] 371 | }, 372 | { 373 | "cell_type": "code", 374 | "execution_count": null, 375 | "id": "17", 376 | "metadata": {}, 377 | "outputs": [], 378 | "source": [ 379 | "def generate_snapshot(images: list):\n", 380 | " \"\"\"Generate a snapshot from images and return an ase atoms.\"\"\"\n", 381 | " image_structs = list(map(AseAtomsAdaptor().get_structure, images))\n", 382 | " sites = set()\n", 383 | " lattice = image_structs[0].lattice\n", 384 | " for site in chain(*(struct for struct in image_structs)):\n", 385 | " sites.add(PeriodicSite(site.species, site.frac_coords, lattice))\n", 386 | " neb_path = Structure.from_sites(sorted(sites))\n", 387 | " return neb_path.to_ase_atoms()" 388 | ] 389 | }, 390 | { 391 | "cell_type": "code", 392 | "execution_count": null, 393 | "id": "18", 394 | "metadata": {}, 395 | "outputs": [], 396 | "source": [ 397 | "# view compressed path b in one snapshot\n", 398 | "view(generate_snapshot(images_b),viewer=\"ngl\")" 399 | ] 400 | }, 401 | { 402 | "cell_type": "code", 403 | "execution_count": null, 404 | "id": "b49b1dc8-356b-4313-98a3-a5101d8fe374", 405 | "metadata": {}, 406 | "outputs": [], 407 | "source": [ 408 | "# view compressed path c in one snapshot\n", 409 | "view(generate_snapshot(images_c),viewer=\"ngl\")" 410 | ] 411 | } 412 | ], 413 | "metadata": { 414 | "kernelspec": { 415 | "display_name": "Python 3 (ipykernel)", 416 | "language": "python", 417 | "name": "python3" 418 | }, 419 | "language_info": { 420 | "codemirror_mode": { 421 | "name": "ipython", 422 | "version": 3 423 | }, 424 | "file_extension": ".py", 425 | "mimetype": "text/x-python", 426 | "name": "python", 427 | "nbconvert_exporter": "python", 428 | "pygments_lexer": "ipython3", 429 | "version": "3.11.0" 430 | }, 431 | "toc": { 432 | "base_numbering": 1, 433 | "nav_menu": {}, 434 | "number_sections": true, 435 | "sideBar": true, 436 | "skip_h1_title": false, 437 | "title_cell": "Table of Contents", 438 | "title_sidebar": "Contents", 439 | "toc_cell": true, 440 | "toc_position": {}, 441 | "toc_section_display": true, 442 | "toc_window_display": false 443 | } 444 | }, 445 | "nbformat": 4, 446 | "nbformat_minor": 5 447 | } 448 | -------------------------------------------------------------------------------- /config/config-01.yml: -------------------------------------------------------------------------------- 1 | 2 | model: MACE 3 | num_interactions: 2 4 | num_channels: 32 5 | max_L: 0 6 | correlations: 2 7 | r_max: 4.0 8 | max_ell: 2 9 | -------------------------------------------------------------------------------- /config/config-02.yml: -------------------------------------------------------------------------------- 1 | 2 | model: "MACE" 3 | num_interactions: 2 4 | num_channels: 32 5 | max_L: 0 6 | correlation: 2 7 | r_max: 4.0 8 | max_ell: 2 9 | name: "mace01" 10 | model_dir: "MACE_models" 11 | log_dir: "MACE_models" 12 | checkpoints_dir: "MACE_models" 13 | results_dir: "MACE_models" 14 | train_file: "data/solvent_xtb_train_200.xyz" 15 | valid_fraction: 0.10 16 | test_file: "data/solvent_xtb_test.xyz" 17 | E0s: "average" 18 | energy_key: "energy_xtb" 19 | forces_key: "forces_xtb" 20 | device: cuda 21 | batch_size: 10 22 | max_num_epochs: 50 23 | swa: True 24 | seed: 123 25 | -------------------------------------------------------------------------------- /data.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Author; alin m elena, alin@elena.re 3 | # Contribs; 4 | # Date: 26-07-2024 5 | # ©alin m elena, GPL v3 https://www.gnu.org/licenses/gpl-3.0.en.html 6 | 7 | from urllib.request import urlretrieve 8 | from pathlib import Path 9 | 10 | default_url = "https://raw.githubusercontent.com/imagdau/Tutorials/main/data/" 11 | 12 | def get_data(url: str=default_url, filename: str| list[str] = "", folder: str="data") -> None: 13 | p = Path(folder) 14 | p.mkdir(parents=True, exist_ok=True) 15 | save_file = p/filename 16 | path, headers = urlretrieve(url+filename, save_file) 17 | if path.exists(): 18 | print(f"saved in {save_file}") 19 | else: 20 | print(f"{save_file} could not be downloaded, check url.") 21 | # example 22 | # from data import get_data 23 | # get_data(filename="LiFePO4_supercell.cif") 24 | -------------------------------------------------------------------------------- /data/LiFePO4_supercell.cif: -------------------------------------------------------------------------------- 1 | # generated using pymatgen 2 | data_LiFePO4 3 | _symmetry_space_group_name_H-M 'P 1' 4 | _cell_length_a 10.23619605 5 | _cell_length_b 11.94151020 6 | _cell_length_c 9.30983438 7 | _cell_angle_alpha 90.0000 8 | _cell_angle_beta 90.00000000 9 | _cell_angle_gamma 90.00000000 10 | _symmetry_Int_Tables_number 1 11 | _chemical_formula_structural LiFePO4 12 | _chemical_formula_sum 'Li16 Fe16 P16 O64' 13 | _cell_volume 1137.99355945 14 | _cell_formula_units_Z 16 15 | loop_ 16 | _symmetry_equiv_pos_site_id 17 | _symmetry_equiv_pos_as_xyz 18 | 1 'x, y, z' 19 | loop_ 20 | _atom_site_type_symbol 21 | _atom_site_label 22 | _atom_site_symmetry_multiplicity 23 | _atom_site_fract_x 24 | _atom_site_fract_y 25 | _atom_site_fract_z 26 | _atom_site_occupancy 27 | Li Li0 1 0.00000000 0.00000000 0.00000000 1 28 | Li Li1 1 0.00000000 0.00000000 0.50000000 1 29 | Li Li2 1 0.00000000 0.50000000 0.00000000 1 30 | Li Li3 1 0.00000000 0.50000000 0.50000000 1 31 | Li Li4 1 0.50000000 0.25000000 0.25000000 1 32 | Li Li5 1 0.50000000 0.25000000 0.75000000 1 33 | Li Li6 1 0.50000000 0.75000000 0.25000000 1 34 | Li Li7 1 0.50000000 0.75000000 0.75000000 1 35 | Li Li8 1 0.50000000 0.00000000 0.25000000 1 36 | Li Li9 1 0.50000000 0.00000000 0.75000000 1 37 | Li Li10 1 0.50000000 0.50000000 0.25000000 1 38 | Li Li11 1 0.50000000 0.50000000 0.75000000 1 39 | Li Li12 1 0.00000000 0.25000000 0.00000000 1 40 | Li Li13 1 0.00000000 0.25000000 0.50000000 1 41 | Li Li14 1 0.00000000 0.75000000 0.00000000 1 42 | Li Li15 1 0.00000000 0.75000000 0.50000000 1 43 | Fe Fe16 1 0.78115127 0.12500000 0.26493287 1 44 | Fe Fe17 1 0.78115127 0.12500000 0.76493287 1 45 | Fe Fe18 1 0.78115127 0.62500000 0.26493287 1 46 | Fe Fe19 1 0.78115127 0.62500000 0.76493287 1 47 | Fe Fe20 1 0.71884873 0.37500000 0.01493386 1 48 | Fe Fe21 1 0.71884873 0.37500000 0.51493386 1 49 | Fe Fe22 1 0.71884873 0.87500000 0.01493387 1 50 | Fe Fe23 1 0.71884873 0.87500000 0.51493386 1 51 | Fe Fe24 1 0.28115127 0.12500000 0.48506664 1 52 | Fe Fe25 1 0.28115127 0.12500000 0.98506664 1 53 | Fe Fe26 1 0.28115127 0.62500000 0.48506663 1 54 | Fe Fe27 1 0.28115127 0.62500000 0.98506664 1 55 | Fe Fe28 1 0.21884873 0.37500000 0.23506663 1 56 | Fe Fe29 1 0.21884873 0.37500000 0.73506664 1 57 | Fe Fe30 1 0.21884873 0.87500000 0.23506663 1 58 | Fe Fe31 1 0.21884873 0.87500000 0.73506664 1 59 | P P32 1 0.09386630 0.12500000 0.20931129 1 60 | P P33 1 0.09386630 0.12500000 0.70931128 1 61 | P P34 1 0.09386630 0.62500000 0.20931129 1 62 | P P35 1 0.09386630 0.62500000 0.70931128 1 63 | P P36 1 0.40613370 0.37500000 0.45931179 1 64 | P P37 1 0.40613370 0.37500000 0.95931179 1 65 | P P38 1 0.40613370 0.87500000 0.45931178 1 66 | P P39 1 0.40613370 0.87500000 0.95931179 1 67 | P P40 1 0.59386630 0.12500000 0.04068772 1 68 | P P41 1 0.59386630 0.12500000 0.54068772 1 69 | P P42 1 0.59386630 0.62500000 0.04068771 1 70 | P P43 1 0.59386630 0.62500000 0.54068772 1 71 | P P44 1 0.90613370 0.37500000 0.29068822 1 72 | P P45 1 0.90613370 0.37500000 0.79068822 1 73 | P P46 1 0.90613370 0.87500000 0.29068822 1 74 | P P47 1 0.90613370 0.87500000 0.79068821 1 75 | O O48 1 0.09423067 0.12500000 0.37239328 1 76 | O O49 1 0.09423067 0.12500000 0.87239328 1 77 | O O50 1 0.09423067 0.62500000 0.37239328 1 78 | O O51 1 0.09423067 0.62500000 0.87239328 1 79 | O O52 1 0.83415452 0.27277902 0.35686767 1 80 | O O53 1 0.83415452 0.27277902 0.85686767 1 81 | O O54 1 0.83415452 0.77277902 0.35686767 1 82 | O O55 1 0.83415452 0.77277902 0.85686767 1 83 | O O56 1 0.83415452 0.47722149 0.35686767 1 84 | O O57 1 0.83415452 0.47722149 0.85686767 1 85 | O O58 1 0.83415452 0.97722149 0.35686767 1 86 | O O59 1 0.83415452 0.97722149 0.85686767 1 87 | O O60 1 0.90577033 0.37500000 0.12760722 1 88 | O O61 1 0.90577033 0.37500000 0.62760722 1 89 | O O62 1 0.90577033 0.87500000 0.12760722 1 90 | O O63 1 0.90577033 0.87500000 0.62760722 1 91 | O O64 1 0.04430856 0.37500000 0.35493222 1 92 | O O65 1 0.04430856 0.37500000 0.85493222 1 93 | O O66 1 0.04430856 0.87500000 0.35493222 1 94 | O O67 1 0.04430856 0.87500000 0.85493222 1 95 | O O68 1 0.40576933 0.37500000 0.12239278 1 96 | O O69 1 0.40576933 0.37500000 0.62239278 1 97 | O O70 1 0.40576933 0.87500000 0.12239278 1 98 | O O71 1 0.40576933 0.87500000 0.62239278 1 99 | O O72 1 0.16584548 0.02277901 0.14313283 1 100 | O O73 1 0.16584548 0.02277902 0.64313283 1 101 | O O74 1 0.16584548 0.52277902 0.14313283 1 102 | O O75 1 0.16584548 0.52277902 0.64313283 1 103 | O O76 1 0.66584548 0.22722098 0.10686717 1 104 | O O77 1 0.66584548 0.22722098 0.60686717 1 105 | O O78 1 0.66584548 0.72722098 0.10686717 1 106 | O O79 1 0.66584548 0.72722098 0.60686717 1 107 | O O80 1 0.33415452 0.27277902 0.39313283 1 108 | O O81 1 0.33415452 0.27277902 0.89313283 1 109 | O O82 1 0.33415452 0.77277902 0.39313283 1 110 | O O83 1 0.33415452 0.77277902 0.89313283 1 111 | O O84 1 0.33415452 0.47722149 0.39313283 1 112 | O O85 1 0.33415452 0.47722149 0.89313283 1 113 | O O86 1 0.33415452 0.97722149 0.39313283 1 114 | O O87 1 0.33415452 0.97722149 0.89313283 1 115 | O O88 1 0.95569144 0.12500000 0.14506827 1 116 | O O89 1 0.95569144 0.12500000 0.64506827 1 117 | O O90 1 0.95569144 0.62500000 0.14506827 1 118 | O O91 1 0.95569144 0.62500000 0.64506828 1 119 | O O92 1 0.45569144 0.12500000 0.10493173 1 120 | O O93 1 0.45569144 0.12500000 0.60493173 1 121 | O O94 1 0.45569144 0.62500000 0.10493173 1 122 | O O95 1 0.45569144 0.62500000 0.60493173 1 123 | O O96 1 0.54430856 0.37500000 0.39506877 1 124 | O O97 1 0.54430856 0.37500000 0.89506878 1 125 | O O98 1 0.54430856 0.87500000 0.39506877 1 126 | O O99 1 0.54430856 0.87500000 0.89506878 1 127 | O O100 1 0.59423067 0.12500000 0.37760772 1 128 | O O101 1 0.59423067 0.12500000 0.87760772 1 129 | O O102 1 0.59423067 0.62500000 0.37760772 1 130 | O O103 1 0.59423067 0.62500000 0.87760772 1 131 | O O104 1 0.66584548 0.02277901 0.10686717 1 132 | O O105 1 0.66584548 0.02277901 0.60686717 1 133 | O O106 1 0.66584548 0.52277902 0.10686717 1 134 | O O107 1 0.66584548 0.52277902 0.60686717 1 135 | O O108 1 0.16584548 0.22722098 0.14313283 1 136 | O O109 1 0.16584548 0.22722099 0.64313283 1 137 | O O110 1 0.16584548 0.72722098 0.14313283 1 138 | O O111 1 0.16584548 0.72722098 0.64313283 1 139 | -------------------------------------------------------------------------------- /data/solvent_liquid.xyz: -------------------------------------------------------------------------------- 1 | 160 2 | Lattice="12.85369 0.0 0.0 0.0 12.85369 0.0 0.0 0.0 12.85369" Properties=species:S:1:pos:R:3:forces:R:3 temperature=600 density=0.9266511910359493 virial="35.33596281215595 9.288776036340266 -7.295002755181301 9.288776036340266 4.136503968853437 -0.8573997863637491 -7.295002755181301 -0.8573997863637491 24.877318409485024" energy=-24026.85743478 pbc="T T T" 3 | C 10.49052625 5.10664893 1.11195987 -1.91391000 0.11637000 -2.91216000 4 | C 9.15657030 5.81709809 1.57492407 -0.92868000 -1.64523000 -1.54407000 5 | O 9.49152461 7.05623236 2.10256805 1.93541000 0.04112000 0.60689000 6 | C 10.83035210 7.04022952 2.08786343 3.87056000 6.43774000 1.26376000 7 | O 11.48704712 8.08374349 2.42527279 -1.84756000 -3.55208000 -0.61670000 8 | O 11.44193067 5.98865629 1.51223663 1.37576000 1.23473000 2.49161000 9 | H 10.73288257 5.07735537 0.00961456 -1.37009000 0.04762000 0.16633000 10 | H 10.55444764 4.07449763 1.32454705 0.50999000 -1.27027000 1.51706000 11 | H 8.58004374 5.32455754 2.36560596 -0.23225000 -0.61513000 -0.52750000 12 | H 8.66406831 6.10673028 0.66109098 -1.32642000 -0.73962000 -0.37425000 13 | C 3.38321975 0.18709189 6.96880798 -5.01777000 1.78895000 -0.37569000 14 | C 2.98364994 1.38817925 6.14537490 2.48471000 0.31689000 -2.27708000 15 | O 1.61516899 1.73333938 6.52711664 3.80320000 1.18480000 -8.09710000 16 | C 1.47668332 1.26331850 7.64139017 -0.72073000 1.45142000 7.56886000 17 | O 0.71543639 1.91865103 8.33485960 -1.48794000 -1.88474000 2.19786000 18 | O 2.47715029 0.44989844 8.11258074 -0.13364000 -1.04694000 -0.26617000 19 | H 4.34231068 0.22517737 7.24783589 3.69929000 -0.22303000 1.30962000 20 | H 3.07507824 -0.62646956 6.41545663 -0.45753000 -2.78184000 -0.54026000 21 | H 3.03248111 1.29004132 5.05337681 -0.36051000 -0.49721000 -0.03341000 22 | H 3.91880731 1.99849815 6.31051910 -1.78646000 1.58440000 0.41753000 23 | C 6.67535824 8.59163134 3.03741692 0.16952000 -1.38286000 2.61316000 24 | C 5.95328935 9.41516727 4.13318114 1.32411000 0.20096000 -0.37986000 25 | O 4.65229026 9.31567971 3.67553836 -3.20461000 3.45769000 1.37442000 26 | C 4.45220972 8.59133571 2.63058478 -0.95036000 -1.25973000 -1.66228000 27 | O 3.40185759 8.35126735 2.05882694 0.37097000 0.45350000 0.53682000 28 | O 5.60789785 8.06822909 2.27511598 1.75323000 -1.67060000 -1.59479000 29 | H 7.35842903 7.82465167 3.48088208 -0.62842000 0.06419000 -0.51340000 30 | H 7.25432700 9.26053738 2.48677770 1.10424000 1.02848000 -1.53304000 31 | H 6.44080410 10.42068573 4.08122653 -0.97839000 -0.05560000 0.24759000 32 | H 5.78850504 9.09116430 5.14441949 1.19559000 -0.56120000 0.84107000 33 | C 5.62985196 11.67898485 10.99304125 0.22493000 -1.62782000 1.46109000 34 | C 4.44305791 10.68657715 11.19453569 1.27465000 4.00940000 5.13924000 35 | O 4.33188633 9.90495713 10.13991614 0.65455000 -0.77042000 -1.58267000 36 | C 5.14341691 10.51846374 9.22000325 -1.53623000 -4.97401000 -6.36010000 37 | O 5.20332796 9.94333824 8.07804287 0.49511000 2.72680000 3.42982000 38 | O 5.68016130 11.63168326 9.60728493 1.42100000 1.61350000 -1.69319000 39 | H 6.55966219 11.22006956 11.42214884 -0.29976000 0.34847000 -0.35675000 40 | H 5.50317884 12.66789634 11.41979661 -0.27726000 0.43836000 -0.23611000 41 | H 3.66835316 11.32270627 11.51860293 -2.44591000 0.79052000 -0.84887000 42 | H 4.69510590 10.39256184 12.17992528 0.31183000 -2.56595000 1.02976000 43 | C 12.56997051 12.02271823 4.54890804 -2.67796000 0.38050000 2.76645000 44 | C 11.61054538 11.78888391 3.40913278 2.19501000 1.91949000 -0.49538000 45 | O 12.24655881 12.59078707 2.38876831 0.61883000 1.70713000 -2.65545000 46 | C 11.84665480 12.44344522 1.07500551 -0.49889000 -3.02301000 0.39599000 47 | O 11.19450998 11.45006063 0.75137530 -0.24248000 0.53284000 -0.25103000 48 | O 12.33022348 13.32718782 0.23188057 -0.39097000 1.91274000 2.00874000 49 | C 12.58155167 13.14248030 -1.14536661 0.43785000 -1.78106000 -0.32205000 50 | H 12.41648459 13.00567847 4.92086812 0.08596000 1.69442000 0.07047000 51 | H 13.63102975 11.89110931 4.62556744 0.69535000 -0.80935000 -1.15130000 52 | H 12.10060516 11.52573030 5.42149364 0.39387000 -0.61238000 -0.09456000 53 | H 10.60394721 12.18833803 3.40534094 -0.50900000 -0.31688000 1.10822000 54 | H 11.70616398 10.78323691 3.22140464 -0.51309000 -2.69303000 -1.26694000 55 | H 12.51205178 14.06193045 -1.68873065 -0.33909000 0.89789000 -0.09278000 56 | H 12.00152892 12.33886759 -1.68068424 0.15263000 0.53470000 0.80711000 57 | H 13.62893461 12.87572766 -1.04764000 0.61404000 -0.27217000 -1.02760000 58 | C 10.39510045 13.19908794 7.09005684 1.75194000 -0.87918000 -0.90114000 59 | C 8.91519085 13.15935719 6.72119450 0.53549000 0.27267000 0.49104000 60 | O 8.87318500 12.64257458 5.38879385 -0.77598000 0.12952000 -2.48503000 61 | C 7.78886055 13.06777463 4.63704579 0.03943000 -1.37616000 0.39140000 62 | O 6.86318922 13.67457163 5.11467606 -0.58332000 0.77694000 -0.00042000 63 | O 7.81454223 12.57025971 3.35856637 1.99019000 -0.87990000 1.33339000 64 | C 6.67665647 12.78773131 2.54032617 -1.18893000 0.91936000 0.93984000 65 | H 10.99824700 13.73279885 6.30672727 -0.86940000 0.05144000 0.21281000 66 | H 11.00399260 12.23603807 7.09514690 -1.25734000 0.42076000 0.25351000 67 | H 10.67491243 13.72157758 7.99201312 -0.41013000 0.26114000 0.76648000 68 | H 8.53610983 14.18497881 6.74560366 -0.12324000 0.20134000 -0.09310000 69 | H 8.19795495 12.56274032 7.31578049 0.51528000 -0.06966000 0.27331000 70 | H 5.76831190 12.76230669 3.18996452 0.22567000 -0.23351000 -0.53486000 71 | H 6.61463742 12.04397823 1.71190585 0.41733000 0.50167000 0.48232000 72 | H 6.78012868 13.84515295 2.26673538 -0.47761000 -0.07940000 -0.90078000 73 | C -0.01885636 4.70457266 9.76014101 -2.17284000 1.14354000 -0.57771000 74 | C 0.61650154 6.01454362 9.29149547 -1.40865000 -1.55631000 -0.99778000 75 | O 1.25562557 6.73148389 10.29917336 0.32920000 1.77805000 1.91791000 76 | C 0.47782308 7.71756757 10.89802677 1.60994000 0.81615000 -0.93554000 77 | O -0.68787807 7.87102777 10.77236910 -3.54263000 -0.98236000 -0.51263000 78 | O 1.20544476 8.66868920 11.46622414 2.03015000 0.83433000 1.01773000 79 | C 0.54500932 9.88112638 12.04401036 -0.37582000 -1.98215000 -1.78747000 80 | H 0.58689948 4.02377697 10.27968716 1.42656000 -0.79665000 0.85833000 81 | H -0.53114018 4.07729973 8.97598880 0.45211000 0.86478000 0.25900000 82 | H -0.77733975 5.09434795 10.43894438 -0.59142000 -0.48986000 0.47032000 83 | H -0.10092717 6.61487521 8.71724402 -0.22444000 0.18945000 0.29865000 84 | H 1.32988134 5.71250761 8.58563509 1.84927000 -0.00871000 -1.13648000 85 | H -0.45104884 9.64420717 12.34080206 -1.20737000 -0.10091000 -0.37139000 86 | H 0.63857133 10.74143670 11.39411494 -0.08140000 -0.10892000 -0.66824000 87 | H 0.97495238 10.11282699 12.94867877 1.87492000 0.24191000 2.20488000 88 | C 5.46937364 4.49190836 1.26859493 0.79957000 1.56307000 0.55169000 89 | C 4.02729245 4.14237796 1.65744476 2.03846000 -0.25695000 -0.91173000 90 | O 4.14720452 2.71624820 1.99854314 5.14115000 3.60283000 4.14241000 91 | C 3.39509656 1.83889389 1.63510505 -8.57040000 -2.65584000 -3.70195000 92 | O 2.45186708 2.14128980 0.84272648 1.20064000 -2.10863000 0.53498000 93 | O 3.44705118 0.65703565 2.35018153 2.23608000 0.93152000 0.83650000 94 | C 2.87164289 -0.58760001 2.07554959 -0.63319000 -1.47435000 -0.82701000 95 | H 6.26992716 3.89626836 1.77754964 -0.74464000 -0.08902000 -0.32592000 96 | H 5.68369607 5.53690050 1.66964292 -0.22619000 -0.72550000 -0.90697000 97 | H 5.61441468 4.44609780 0.17217518 0.15818000 0.07123000 0.37381000 98 | H 3.40039228 4.23028435 0.78924227 -0.68321000 -0.00627000 -0.32476000 99 | H 3.73201747 4.86060360 2.43339632 -0.51758000 -1.01347000 0.65909000 100 | H 1.84319345 -0.50383251 1.74483700 -0.44814000 0.13278000 -0.05147000 101 | H 2.94701693 -1.33128880 2.88606472 -0.29666000 1.08459000 1.15911000 102 | H 3.38883681 -1.30347342 1.41525554 0.56890000 0.92370000 -1.19830000 103 | C 2.14384125 8.28721740 6.28181682 -0.11078000 -3.01958000 1.15412000 104 | C 2.07695065 7.44429812 4.98121619 -0.64787000 -0.45772000 3.41039000 105 | O 2.64725602 6.11189747 5.25344449 0.63816000 -0.88896000 0.50505000 106 | C 3.98373129 6.04491689 5.35470586 2.87260000 0.66204000 -1.58091000 107 | O 4.83648365 6.93259272 5.18030700 -1.94056000 -0.28421000 0.39263000 108 | O 4.46970361 4.75774837 5.49149483 -1.61923000 -0.77087000 -0.11683000 109 | C 5.80880103 4.44548084 5.08378864 2.03502000 0.85601000 3.77097000 110 | H 3.05338121 8.35146013 6.85692947 1.88775000 -0.03287000 -0.96989000 111 | H 1.73673919 7.61840135 7.08026233 -1.49850000 1.01199000 -0.31137000 112 | H 1.89007085 9.30621938 6.20177689 -1.41644000 1.18687000 -0.12033000 113 | H 2.73264309 7.88318736 4.35572993 1.67204000 1.47587000 -2.99178000 114 | H 1.04476078 7.27058050 4.61667269 0.56165000 0.10690000 0.10819000 115 | H 6.25729198 3.67807697 5.80909666 -1.34122000 -0.14946000 -1.23847000 116 | H 6.54790106 5.21167644 5.45382067 -0.79204000 0.44079000 -0.80054000 117 | H 6.00586095 4.25804831 4.05398956 -0.44838000 -0.36660000 -1.05550000 118 | C 8.13891795 4.54564963 8.61156098 0.68313000 -1.41454000 0.72964000 119 | C 8.70093270 4.93917821 9.97840952 1.66515000 -0.82046000 -0.85278000 120 | O 8.44409026 4.29839605 11.24120744 -1.68945000 -1.50251000 2.09538000 121 | C 9.01582238 3.11354291 11.59761456 1.19901000 -1.80369000 0.55702000 122 | O 9.72044882 2.39237378 10.92584216 0.12078000 0.26162000 -1.42080000 123 | O 8.83417405 2.72061846 12.91560622 -3.26106000 4.17375000 0.01667000 124 | C 9.22909082 1.54662668 13.57316244 -0.04726000 -1.11245000 -2.57979000 125 | H 7.18512274 4.09084752 8.79302938 -0.79401000 -0.82939000 -0.27824000 126 | H 8.83368561 3.94346711 7.96783533 -0.55852000 -0.04320000 0.78075000 127 | H 7.92174201 5.40860782 8.03918617 -0.19880000 1.43171000 -0.47176000 128 | H 8.53126399 5.96100800 10.12766657 -0.75759000 1.56572000 0.70440000 129 | H 9.74911255 4.85228726 9.96743247 1.90022000 0.10264000 -0.72137000 130 | H 10.25692327 1.72166823 13.66947514 2.09486000 -0.85727000 0.28189000 131 | H 8.94501141 0.67045490 12.93893567 0.67412000 0.20096000 0.30965000 132 | H 8.81630742 1.37095530 14.54242065 -0.91302000 0.69092000 0.79398000 133 | C 11.29694105 3.74539175 5.38853677 -0.97567000 0.48743000 -0.41402000 134 | C 11.60022387 5.30401734 5.45352503 -0.45083000 -2.17684000 0.53318000 135 | O 10.37809502 6.01816836 5.63962076 -3.13240000 0.39236000 -1.36095000 136 | C 10.44163081 7.27797137 6.14193011 4.26565000 -0.74326000 0.26852000 137 | O 11.46391048 7.67499615 6.63857098 2.32610000 0.24489000 1.40238000 138 | O 9.48390237 8.08158407 5.98114330 -4.93447000 1.89889000 -1.43800000 139 | C 9.47280962 9.42259670 6.61138543 -3.08622000 -5.90139000 4.73302000 140 | H 10.34671916 3.74202408 4.84138090 -0.38560000 -0.43896000 -0.07869000 141 | H 11.11965010 3.17718153 6.35785925 0.42732000 1.03033000 -0.74232000 142 | H 11.96996026 3.14819646 4.74213756 0.23875000 0.47942000 0.36075000 143 | H 12.17433393 5.64663245 4.58866450 0.01851000 0.22965000 0.03706000 144 | H 12.01203039 5.44128190 6.39761571 2.25193000 0.48497000 1.40234000 145 | H 9.90178868 10.16385613 6.18582546 4.15527000 4.37943000 -3.44865000 146 | H 9.68628371 9.29537087 7.70379483 0.71461000 0.08333000 -0.78295000 147 | H 8.46328082 9.77086742 6.69205519 -1.42585000 -0.46970000 -0.55013000 148 | C 4.57065648 5.87213759 9.42433836 0.45167000 -0.79642000 -1.27392000 149 | C 5.35276495 6.65613556 10.51712052 -0.04280000 -3.20255000 -0.84715000 150 | O 6.42992988 7.36312707 9.88767532 1.69484000 2.76021000 0.09364000 151 | C 7.15570063 8.18722854 10.76244605 -1.79247000 -1.13023000 4.43817000 152 | O 6.77925462 8.44062619 11.94947147 -0.56013000 -0.25626000 -2.57522000 153 | O 8.40040057 8.49353197 10.47282671 0.18948000 -0.48785000 -2.63998000 154 | C 9.03502580 9.50315077 11.30599004 2.26568000 -1.14724000 -2.32500000 155 | H 3.72082191 5.13108380 9.61599973 0.50157000 1.64786000 0.79401000 156 | H 5.21363661 5.13136658 8.89458638 0.45298000 0.57692000 0.17073000 157 | H 4.31304282 6.60177730 8.64347669 -0.51300000 -0.09608000 -0.00631000 158 | H 5.94345628 5.91506891 11.04257937 -0.34716000 -0.47429000 1.10209000 159 | H 4.64560634 7.13361158 11.14977915 -0.18022000 1.37172000 0.65728000 160 | H 8.36349763 10.30024665 11.36655663 -1.06322000 1.37050000 1.03671000 161 | H 9.88077289 9.84527743 10.62671108 -0.60030000 -0.36805000 0.84030000 162 | H 9.47137002 9.02067467 12.17632624 -0.30994000 0.31461000 0.54801000 163 | -------------------------------------------------------------------------------- /figures/active_learning.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/active_learning.png -------------------------------------------------------------------------------- /figures/embedding_layer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/embedding_layer.png -------------------------------------------------------------------------------- /figures/interaction_first_layer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/interaction_first_layer.png -------------------------------------------------------------------------------- /figures/interaction_layer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/interaction_layer.png -------------------------------------------------------------------------------- /figures/interaction_layer_second.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/interaction_layer_second.png -------------------------------------------------------------------------------- /figures/iterative_training.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/iterative_training.png -------------------------------------------------------------------------------- /figures/mace_architecture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/mace_architecture.png -------------------------------------------------------------------------------- /figures/mace_architecture_layers.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/mace_architecture_layers.png -------------------------------------------------------------------------------- /figures/product_layer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/product_layer.png -------------------------------------------------------------------------------- /figures/readout.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/readout.png -------------------------------------------------------------------------------- /figures/sketch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/imagdau/Tutorials/3b7b691f60afdffc0cd66948e333883ae1689cd8/figures/sketch.png --------------------------------------------------------------------------------