def energy(R, **kwargs):
dr = metric(R, R, **kwargs)
total_charge = util.high_precision_sum(charge_fn(dr), axis=1)
embedding_energy = embedding_fn(total_charge)
pairwise_energy = util.high_precision_sum(
smap._diagonal_mask(pairwise_fn(dr)), axis=1) / f32(2.0)
return util.high_precision_sum(embedding_energy + pairwise_energy,
axis=axis)
def energy(R, **kwargs):
dr = space.distance(displacement(R, R, **kwargs))
total_charge = smap._high_precision_sum(charge_fn(dr), axis=1)
embedding_energy = embedding_fn(total_charge)
pairwise_energy = smap._high_precision_sum(
smap._diagonal_mask(pairwise_fn(dr)), axis=1) / f32(2.0)
return smap._high_precision_sum(embedding_energy + pairwise_energy,
axis=axis)