Skip to content

Commit

Permalink
total charge
Browse files Browse the repository at this point in the history
  • Loading branch information
yuanqing-wang committed Aug 25, 2021
1 parent d7c896b commit 489ab27
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions espaloma/nn/readout/charge_equilibrium.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ class ChargeEquilibrium(torch.nn.Module):
def __init__(self):
super(ChargeEquilibrium, self).__init__()

def forward(self, g):
def forward(self, g, total_charge=0.0):
""" apply charge equilibrium to all molecules in batch """
# calculate $s ^ {-1}$ and $ es ^ {-1}$
g.apply_nodes(lambda node: {
Expand All @@ -78,9 +78,10 @@ def forward(self, g):
dgl.function.sum(msg='m_q', out='sum_q'),
etype='n1_in_g')
else:
g.nodes['g'].data['sum_q'] = torch.zeros(
g.nodes['g'].data['sum_q'] = torch.ones(
g.batch_size, 1,
)
device=g.nodes['n1'].data['s'].device,
) * total_charge

g.update_all(
dgl.function.copy_src(src='sum_q', out='m_sum_q'),
Expand Down

0 comments on commit 489ab27

Please sign in to comment.