Appendix#
A1: Using layers from a TorchScript model#
import torch
import torch.nn as nn
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 1)
def forward(self, x):
# do nothing
return x
model = MyModel()
scripted_model = torch.jit.script(model)
x = torch.rand(10)
print(scripted_model(x))
# Output: tensor([0.7581, 0.3445, 0.1828, 0.5155, 0.1966, 0.6062, 0.5789, 0.8587, 0.5662,
# 0.3897])
# =========================================
# Wrapped model to use layers explicitly
# =========================================
class WrapperModel(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.module_list = torch.nn.ModuleList(list(model.children()))
def forward(self, x):
x = self.module_list[0](x)
x = self.module_list[1](x)
return x
wrapped_model = WrapperModel(scripted_model)
scripted_wrapped_model = torch.jit.script(wrapped_model)
print(scripted_wrapped_model(x))
# Output: tensor([-0.5165], grad_fn=<AddBackward0>)
A2: Pytorch implementation of Stillinger-Weber energy functions#
def calc_sw2(A, B, p, q, sigma, cutoff, rij):
if rij < cutoff:
sig_r = sigma / rij
E2 = A * (B * sig_r ** p - sig_r ** q) * torch.exp(sigma / (rij - cutoff))
else:
E2 = torch.tensor(0.0, dtype=torch.float32)
return E2
def calc_sw3(
lam, cos_beta0, gamma_ij, gamma_ik, cutoff_ij, cutoff_ik, cutoff_jk, rij, rik, rjk
):
cos_beta_ikj = (rij ** 2 + rik ** 2 - rjk ** 2) / (2 * rij * rik)
cos_diff = cos_beta_ikj - cos_beta0
exp_ij_ik = torch.exp(gamma_ij / (rij - cutoff_ij) + gamma_ik / (rik - cutoff_ik))
E3 = lam * exp_ij_ik * cos_diff ** 2
return E3
def energy(particle_contributing: torch.Tensor, coords: torch.Tensor,
num_neighbors: torch.Tensor, neighbor_list: torch.Tensor, A: torch.Tensor,
B: torch.Tensor, p: torch.Tensor, q: torch.Tensor, sigma: torch.Tensor,
gamma: torch.Tensor, cutoff: torch.Tensor, lam: torch.Tensor, cos_beta0: torch.Tensor,
):
energy_conf = torch.tensor(0.0, dtype=torch.float32)
neigh_list_cursor = 0
num_atoms = particle_contributing.shape[0]
for atom_i in range(num_atoms):
if particle_contributing[atom_i] != 1:
continue
# Coordinates of atom i
xyz_i = coords[0, atom_i * 3 : (atom_i + 1) * 3]
num_neigh_i = int(num_neighbors[atom_i])
neigh_i_begin = neigh_list_cursor
neigh_i_end = neigh_i_begin + num_neigh_i
nli = neighbor_list[0, torch.arange(neigh_i_begin, neigh_i_end)]
neigh_list_cursor = neigh_i_end
for atom_j in range(num_neigh_i):
xyz_j = coords[0, (nli[atom_j]) * 3 : ((nli[atom_j]) + 1) * 3]
rij = xyz_j - xyz_i
norm_rij = (rij[0] ** 2 + rij[1] ** 2 + rij[2] ** 2) ** 0.5
E2 = calc_sw2(A, B, p, q, sigma, cutoff, norm_rij)
energy_conf = energy_conf + 0.5 * E2
for atom_k in range(atom_j + 1, num_neigh_i):
xyz_k = coords[0, (nli[atom_k]) * 3 : ((nli[atom_k]) + 1) * 3]
rik = xyz_k - xyz_i
norm_rik = (rik[0] ** 2 + rik[1] ** 2 + rik[2] ** 2) ** 0.5
rjk = xyz_k - xyz_j
norm_rjk = (rjk[0] ** 2 + rjk[1] ** 2 + rjk[2] ** 2) ** 0.5
E3 = calc_sw3( lam, cos_beta0, gamma, gamma, cutoff, cutoff, cutoff,
norm_rij, norm_rik, norm_rjk)
energy_conf = energy_conf + E3
return energy_conf
A3: Example descriptor.dat file#
TorchML model driver descriptor file for Behler Symmetry funcitons (KLIFF set 51).
#================================================================================
# Descriptor parameters file generated by KLIFF.
#================================================================================
cos # cutoff type
1 # number of species
# species 1 species 2 cutoff
Si Si 3.0
#================================================================================
# symmetry functions
#================================================================================
2 # number of symmetry functions types
# sym_function rows cols
g2 8 2
0.0035710676725828126 0.0 # eta Rs
0.03571067672582813 0.0 # eta Rs
0.07142135345165626 0.0 # eta Rs
0.12498736854039845 0.0 # eta Rs
0.21426406035496876 0.0 # eta Rs
0.3571067672582813 0.0 # eta Rs
0.7142135345165626 0.0 # eta Rs
1.428427069033125 0.0 # eta Rs
g4 43 3
1 -1 0.00035710676725828126 # zeta lambda eta
1 1 0.00035710676725828126 # zeta lambda eta
2 -1 0.00035710676725828126 # zeta lambda eta
2 1 0.00035710676725828126 # zeta lambda eta
1 -1 0.010713203017748437 # zeta lambda eta
1 1 0.010713203017748437 # zeta lambda eta
2 -1 0.010713203017748437 # zeta lambda eta
2 1 0.010713203017748437 # zeta lambda eta
1 -1 0.0285685413806625 # zeta lambda eta
1 1 0.0285685413806625 # zeta lambda eta
2 -1 0.0285685413806625 # zeta lambda eta
2 1 0.0285685413806625 # zeta lambda eta
1 -1 0.05356601508874219 # zeta lambda eta
1 1 0.05356601508874219 # zeta lambda eta
2 -1 0.05356601508874219 # zeta lambda eta
2 1 0.05356601508874219 # zeta lambda eta
4 -1 0.05356601508874219 # zeta lambda eta
4 1 0.05356601508874219 # zeta lambda eta
16 -1 0.05356601508874219 # zeta lambda eta
16 1 0.05356601508874219 # zeta lambda eta
1 -1 0.08927669181457032 # zeta lambda eta
1 1 0.08927669181457032 # zeta lambda eta
2 -1 0.08927669181457032 # zeta lambda eta
2 1 0.08927669181457032 # zeta lambda eta
4 -1 0.08927669181457032 # zeta lambda eta
4 1 0.08927669181457032 # zeta lambda eta
16 -1 0.08927669181457032 # zeta lambda eta
16 1 0.08927669181457032 # zeta lambda eta
1 -1 0.16069804526622655 # zeta lambda eta
1 1 0.16069804526622655 # zeta lambda eta
2 -1 0.16069804526622655 # zeta lambda eta
2 1 0.16069804526622655 # zeta lambda eta
4 -1 0.16069804526622655 # zeta lambda eta
4 1 0.16069804526622655 # zeta lambda eta
16 -1 0.16069804526622655 # zeta lambda eta
16 1 0.16069804526622655 # zeta lambda eta
1 -1 0.28568541380662504 # zeta lambda eta
1 1 0.28568541380662504 # zeta lambda eta
2 -1 0.28568541380662504 # zeta lambda eta
2 1 0.28568541380662504 # zeta lambda eta
4 -1 0.28568541380662504 # zeta lambda eta
4 1 0.28568541380662504 # zeta lambda eta
16 1 0.28568541380662504 # zeta lambda eta
#================================================================================
# Preprocessing data to center and normalize
#================================================================================
center_and_normalize False
51 # descriptor size
# mean
0.0
# standard deviation
1.0
Caution
TorchML model driver ignores the center_and_normalize, hence please avoid it for now.