From d5597f9c96db5b38cf19559c7d53641a71ae7a29 Mon Sep 17 00:00:00 2001 From: ekindogus Date: Mon, 8 Apr 2024 14:38:34 -0700 Subject: [PATCH] Update electrostatics.py fix bug --- jax_md/_energy/electrostatics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jax_md/_energy/electrostatics.py b/jax_md/_energy/electrostatics.py index 619ba040..2049f4c4 100644 --- a/jax_md/_energy/electrostatics.py +++ b/jax_md/_energy/electrostatics.py @@ -129,7 +129,7 @@ def energy_fn(position, **kwargs): S2 = jnp.abs(structure_factor(g, position, charge))**2 fn = lambda g2: jnp.exp(-g2 / (4*alpha**2)) / g2 * S2 - return Z * util.high_precision_sum(safe_mask(mask, fn, g2, 1)) + return Z * util.high_precision_sum(util.safe_mask(mask, fn, g2, 1)) return energy_fn