Espaloma trained on QCArchive-derived bond/angle parameters

[bieniekmateusz]

We have derived molecule-specific bond/angle force field parameters for ~20K molecules hosted on QCArchive.

We would like to use these data to train espaloma to assign bond/angle parameters to any organic molecule outside the training set.

Please can you let us know if this sounds possible, and if so how we can get started with the code?

@djcole56
@jthorton

[yuanqing-wang]

Sounds exciting! This is definitely doable.

It could be done as:

import espaloma as esp

# define a layer
layer = esp.nn.layers.dgl_legacy.gn("GraphConv")

# define a representation
representation = esp.nn.Sequential(
        layer,
        [32, 'relu', 32, 'relu', 32, 'relu']
)

# define a readout
readout = esp.nn.readout.janossy.JanossyPooling(
        in_features=32,
        out_features={
            2: ["k", "eq"],
        },
) 

net = torch.nn.Sequential(
    representation,
    readout
)

and then suppose you have a molecule, say methane, and you know that the reference force constant and equilibrium bond length, and you want the loss function to be simply the sum of the MSE loss:

g = esp.Graph("C")
g = net(g)
k_hat = g.nodes['n2'].data['k']
eq_hat = g.nodes['n2'].data['eq']

loss = torch.nn.MSELoss()(k_ref, k_hat) + torch.nn.MSELoss()(eq_ref, eq_hat)
loss.backward()

you can also look at this pre-defined loss function

espaloma/espaloma/metrics.py

Line 417 in e0797bb

class BondKMSE(GraphMetric):

[bieniekmateusz]

Thanks @yuanqing-wang, would you mind double checking if I am not doing anything silly?

I was just trying to run the example you shared here, and run into this:

  File "python3.9/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "python3.9/site-packages/torch/nn/modules/container.py", line 119, in forward
    input = module(input)
  File "python3.9/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "python3.9/site-packages/espaloma/nn/sequential.py", line 119, in forward
    g_ = dgl.to_homo(g.edge_type_subgraph(["n1_neighbors_n1"]))
AttributeError: 'Graph' object has no attribute 'edge_type_subgraph'

Was that meant to be g_ = dgl.to_homo(g.heterograph.edge_type_subgraph(["n1_neighbors_n1"]))? It shows up in a ~3 other places.

[yuanqing-wang]

Hi sorry for the late reply! If the graph you're having is a espaloma.Graph, which is a thin wrapper of dgl.graph then you would indeed need to convert it to dgl.graph by using g.heterograph!