From d7b465ef79412e730ba1c3c471335cdb8e6efd6c Mon Sep 17 00:00:00 2001 From: Marcel Langer Date: Fri, 18 Oct 2024 13:36:37 +0200 Subject: [PATCH] Remove for loop from EwaldCalculator Fix #80, suggested by @ceriottm --- src/torchpme/calculators/ewald.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/src/torchpme/calculators/ewald.py b/src/torchpme/calculators/ewald.py index 75e3adda..1f1e0e68 100644 --- a/src/torchpme/calculators/ewald.py +++ b/src/torchpme/calculators/ewald.py @@ -117,18 +117,14 @@ def _compute_kspace( charges_reshaped = charges sum_idx = 1 + # Actual computation of trigonometric factors # Actual computation of trigonometric factors cos_all = torch.cos(trig_args) sin_all = torch.sin(trig_args) - cos_summed = torch.sum(cos_all * charges_reshaped, dim=sum_idx) - sin_summed = torch.sum(sin_all * charges_reshaped, dim=sum_idx) - - # Add up the contributions to compute the potential - energy = torch.zeros_like(charges) - for i in range(num_atoms): - energy[i] += torch.sum( - G * cos_all[:, i] * cos_summed, dim=sum_idx - 1 - ) + torch.sum(G * sin_all[:, i] * sin_summed, dim=sum_idx - 1) + cos_summed_G = torch.sum(cos_all * charges_reshaped, dim=sum_idx) * G + sin_summed_G = torch.sum(sin_all * charges_reshaped, dim=sum_idx) * G + + energy = (cos_summed_G @ cos_all + sin_summed_G @ sin_all).T energy /= torch.abs(cell.det()) # Remove the self-contribution: Using the Coulomb potential as an