Integrate PBC to the config by APerezFadon · Pull Request #112 · google-deepmind/ferminet

APerezFadon · 2026-05-25T15:04:33Z

Periodic boundary conditions can now be imposed by setting the cfg.system.lattice option to the config. If set, evaluate the Coulomb interaction with an Ewald sum (in 2D or 3D), use periodic envelopes, input features, and Jastrow factor. After each MCMC step, also map electrons back to the simulation cell.

jsspencer

Thanks for this! Looks really useful, but I think we need a bit of cleaning up to avoid a few issues with existing code.

jsspencer · 2026-05-25T16:57:38Z

@@ -0,0 +1,146 @@
+# Copyright 2022 DeepMind Technologies Limited.


Can you update the year please?

jsspencer · 2026-05-25T16:58:08Z

@@ -0,0 +1,147 @@
+# Copyright 2022 DeepMind Technologies Limited.


Also year here

jsspencer · 2026-05-25T16:59:36Z

    """
    del key  # unused
-    potential_energy = make_ewald_potential(
+    potential_energy = make_potential_energy(


The previous name was more descriptive. Is there a reason not to keep it?

Sure, reverted to the original name

jsspencer · 2026-05-25T17:01:01Z



-def make_simple_ee_jastrow():
+def make_simple_ee_jastrow(lattice: Optional[jnp.ndarray] = None, ndim: int = 3):


Nit: jnp.ndarray | None is the modern way

Done (in some other places too)

jsspencer · 2026-10-05T16:04:59Z

+  madelung_const += (
+      (2 * jnp.pi / volume) *
+      jnp.sum(jax.scipy.special.erfc(rec_vec_norm / (2 * root_gamma)) / rec_vec_norm) -
+      2 * (jnp.pi**0.5) / (volume * root_gamma))


At a glance: ewald2d.py and ewald3d.py look identical apart from the Madelung constant. Can we deduplicate them and handle the two cases in the same function?

There are some more differences, but the 2D and 3D Ewalds have been combined into a single function

jsspencer · 2026-10-05T18:22:52Z

+  cfg.system.molecule = [system.Atom('X', (0, 0, 0))]
+
+  cfg.system.lattice = _sc_lattice_vecs(rs, sum(cfg.system.electrons), cfg.system.ndim)
+  cfg.system.make_local_energy_kwargs['heg'] = True


This feels a little error prone -- I wonder if it's simpler to have a custom make_local_energy_function just for the HEG (similar to previous version) which just hardcodes this flag. What do you think?

The total charge in a simulation cell should always be zero, otherwise the potential always diverges. Now the ewald potential claculates the background charge of all particles always (including the atoms). Therfore, if there are atoms present and the cell is already charge neutral then nothing happens (the backgrounds cancel). If we are doing the HEG, then the backgrounds do not cancel, and is equivalent to the old code with include_heg_background = True.

jsspencer · 2026-10-05T18:42:29Z

+    norm = make_periodic_r_ee(lattice)
+  else:
+    norm = lambda x: jnp.linalg.norm(x, axis = -1, keepdims = True)
+


Some care is needed to avoid NaN in the gradient here (though we should not explicitly compute the diagonal terms):

ferminet/ferminet/networks.py

Line 449 in c4312c3

def construct_input_features(

and similarly in the periodic norm code.

How about instead we pass either r_ee or the periodic version through to the jastrow? Something like in networks.py:

if jastrow_apply is not None: if options.lattice is not None: eye = jnp.eye(ee.shape[0]) s_ee = jnp.einsum('il,jkl->jki', reciprocal_vecs, ee) + eye[..., None] jastrow_r_ee = (periodic_norm(lattice_metric, s_ee) * (1.0 - eye))[..., None] else: jastrow_r_ee = r_ee jastrow = jnp.exp( jastrow_apply(jastrow_r_ee, params['jastrow'], nspins) / sum(nspins) )

A bit more work would be to have make_orbitals re-use the computation of r_ae and r_ee to be periodic versions, but that's outside scope here.

The periodic norm was already "safe", the s_ee += jnp.eye(n)[..., None] line makes it so that derivatives of |0| are not computed, which is the source of NaNs. This wasn't done in the Jastrow for open boundary conditions, but now it is. I have not been able to produce any NaNs in my testing.

jsspencer · 2026-10-05T18:49:40Z

Do we also need to ensure the initial (random) electron positions are mapped within the simulation cell?

Once you do a single MCMC step (including burn-in), the walkers are mapped to the simulation cell. As this happens before any energies or other observables are computed, I don't think it is needed.

Only if the proposed move in MCMC is accepted.

Initial positions are now also mapped to the simulation cell

jsspencer · 2026-10-05T18:52:08Z

 def mh_update(
    params: networks.ParamTree,
    f: networks.LogFermiNetLike,
+    _map: Callable,


Nit: this is a bit of an unclear name and is placed high in the argument list.

How about having it as an optional argument, with default for OBC case:

map_to_cell: Callable = lambda x: x

or similar?

jsspencer · 2026-10-05T18:55:47Z

+    x2 = _map(x2)
    lp_2 = 2.0 * f(
        params, x2, data.spins, data.atoms, data.charges
    )  # log prob of proposal


This was preivously only working by accident (f handled the 4D case) but _map flattens x2. We should reshape x2 here to be 4D to match the shape of x1.

…ariables renamed, and safe norm in the Jastrow

jsspencer · 2026-10-06T21:37:02Z

 from ferminet import networks
-from ferminet.pbc.ewald2d import make_ewald_potential as ewald2d
-from ferminet.pbc.ewald3d import make_ewald_potential as ewald3d
+from ferminet.pbc.ewald import make_ewald_potential


Can you match the existing style (import modules rather than directly importing functions and classes into the module namespace) please?

from ferminet.pbc import ewald ... potential_energy = ewald.make_ewald_potential(...)

jsspencer · 2026-10-06T21:46:21Z

+def safe_norm(ee: jnp.ndarray):
+  n = ee.shape[0]
+  return (
+      jnp.linalg.norm(ee + jnp.eye(n)[..., None], axis=-1) * (1.0 - jnp.eye(n)))


Note minor inconsistency here: safe_norm returns shape (nelec, nelec) but make_periodic_r_ee returns a function which returns shape (nelec, nelec, 1)

safe_norm now returns shape (nelec, nelec, 1) for consistency with the previous Jastrow code.

jsspencer · 2026-10-06T21:53:35Z

+  atom_atom_potential += 0.5 * jnp.sum(charges**2) * madelung_const

  def potential(ae: jnp.ndarray, ee: jnp.ndarray):
    """Accumulates atom-electron, atom-atom, and electron-electron potential."""


Do the tests still pass? The original Ewald code reduced vectors to the unit cell to ensure the potential was periodic and this is tested in pbc/tests/hamiltonian_test.py.

The potential now maps positions to the simulation cell and tests pass.

…cell

jsspencer

Thanks!

-- 985c1d8 by Andres <andresperezfadon@gmail.com>: Integrate PBC to the config -- 79269c7 by Andres <andresperezfadon@gmail.com>: Address PR comments: Ewald potentials combined and refactored, MCMC variables renamed, and safe norm in the Jastrow -- a5e3e8c by Andres <andresperezfadon@gmail.com>: Map inital pos, and walkers in the Ewald potential to the simulation cell COPYBARA_INTEGRATE_REVIEW=#112 from APerezFadon:main a5e3e8c PiperOrigin-RevId: 996364455 Change-Id: Ie5dc925d828289eea347eeefd5f6d57ad674f4b5

jsspencer · 2026-10-15T11:03:15Z

Merged in 91cf270. Thanks!

Integrate PBC to the config

985c1d8

jsspencer reviewed Oct 5, 2026

View reviewed changes

Address PR comments: Ewald potentials combined and refactored, MCMC v…

79269c7

…ariables renamed, and safe norm in the Jastrow

jsspencer reviewed Oct 6, 2026

View reviewed changes

Map inital pos, and walkers in the Ewald potential to the simulation …

a5e3e8c

…cell

jsspencer approved these changes Oct 7, 2026

View reviewed changes

jsspencer closed this Oct 15, 2026

		@@ -0,0 +1,146 @@
		# Copyright 2022 DeepMind Technologies Limited.

		@@ -0,0 +1,147 @@
		# Copyright 2022 DeepMind Technologies Limited.



		def make_simple_ee_jastrow():
		def make_simple_ee_jastrow(lattice: Optional[jnp.ndarray] = None, ndim: int = 3):

Conversation

APerezFadon commented May 25, 2026

Uh oh!

jsspencer left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jsspencer left a comment

Choose a reason for hiding this comment

Uh oh!

jsspencer commented Oct 15, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants