Integrate PBC to the config

ferminet/base_config.py

@@ -148,10 +148,12 @@ def default() -> ml_collections.ConfigDict:
           # 1. Specify the system by setting variables below.
           # list of system.Atom objects with element type and position.
           'molecule': config_dict.placeholder(list),
+          # If None, assume OBC, if PBC then matrix with the 
+          # supercell lattice vectors
+          'lattice': None,
           # number of spin up, spin-down electrons
           'electrons': tuple(),
-          # Dimensionality. Change with care. FermiNet implementation currently
-          # assumes 3D systems.
+          # Dimensionality.
           'ndim': 3,
           # Number of excited states. If 0, use normal ground state machinery.
           # If 1, compute ground state using excited state machinery. If >1,

ferminet/configs/heg.py

@@ -15,43 +15,41 @@
 """Unpolarised 14 electron simple cubic homogeneous electron gas."""
 
 from ferminet import base_config
-from ferminet.pbc import envelopes
 from ferminet.utils import system
 
 import numpy as np
 
 
-def _sc_lattice_vecs(rs: float, nelec: int) -> np.ndarray:
+def _sc_lattice_vecs(rs: float, nelec: int, ndim: int) -> np.ndarray:
   """Returns simple cubic lattice vectors with Wigner-Seitz radius rs."""
-  volume = (4 / 3) * np.pi * (rs**3) * nelec
-  length = volume**(1 / 3)
-  return length * np.eye(3)
-
+  if ndim == 2:
+    area = np.pi * (rs**2) * nelec
+    length = area**(1 / 2)
+    return length * np.eye(2)
+  elif ndim == 3:
+    area = 4 * np.pi * (rs**3) * nelec / 3
+    length = area**(1 / 3)
+    return length * np.eye(3)
+  else:
+    raise NotImplementedError
+
+rs = 1.0
 
 def get_config():
-  """Returns config for running unpolarised 14 electron gas with FermiNet."""
   # Get default options.
   cfg = base_config.default()
+
+  # SYSTEM
   cfg.system.electrons = (7, 7)
   # A ghost atom at the origin defines one-electron coordinate system.
   # Element 'X' is a dummy nucleus with zero charge
-  cfg.system.molecule = [system.Atom("X", (0., 0., 0.))]
+  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
+  cfg.network.make_feature_layer_kwargs['include_r_ae'] = False
+
   # Pretraining is not currently implemented for systems in PBC
   cfg.pretrain.method = None
 
-  lattice = _sc_lattice_vecs(1.0, sum(cfg.system.electrons))
-  kpoints = envelopes.make_kpoints(lattice, cfg.system.electrons)
-
-  cfg.system.make_local_energy_fn = "ferminet.pbc.hamiltonian.local_energy"
-  cfg.system.make_local_energy_kwargs = {"lattice": lattice, "heg": True}
-  cfg.network.make_feature_layer_fn = (
-      "ferminet.pbc.feature_layer.make_pbc_feature_layer")
-  cfg.network.make_feature_layer_kwargs = {
-      "lattice": lattice,
-      "include_r_ae": False
-  }
-  cfg.network.make_envelope_fn = (
-      "ferminet.pbc.envelopes.make_multiwave_envelope")
-  cfg.network.make_envelope_kwargs = {"kpoints": kpoints}
-  cfg.network.full_det = True
   return cfg

ferminet/jastrows.py

@@ -15,7 +15,7 @@
 """Multiplicative Jastrow factors."""
 
 import enum
-from typing import Any, Callable, Iterable, Mapping, Union
+from typing import Any, Callable, Iterable, Mapping, Union, Optional
 
 import jax.numpy as jnp
 
@@ -29,10 +29,42 @@ class JastrowType(enum.Enum):
   SIMPLE_EE = enum.auto()
 
 
+def make_periodic_r_ee(lattice: jnp.ndarray):
+  """Util function for transforming r_ee into the periodic version.
+  Some of this could be refactored with ferminet.pbc.feature_layer
+
+  Args:
+      lattice: Matrix whose columns are the primitive lattice vectors of the
+        system, shape (ndim, ndim).
+  """
+
+  # Calculate reciprocal vectors, factor 2pi omitted
+  reciprocal_vecs = jnp.linalg.inv(lattice)
+  lattice_metric = lattice.T @ lattice
+
+  def apply(ee: jnp.ndarray):
+    s_ee = jnp.einsum('il,jkl->jki', reciprocal_vecs, ee)
+
+    n = ee.shape[0]
+    s_ee += jnp.eye(n)[..., None]
+
+    a = (1 - jnp.cos(2 * jnp.pi * s_ee))
+    b = jnp.sin(2 * jnp.pi * s_ee)
+    cos_term = jnp.einsum('...m,mn,...n->...', a, lattice_metric, a)
+    sin_term = jnp.einsum('...m,mn,...n->...', b, lattice_metric, b)
+    periodic_r_ee = (1 / (2 * jnp.pi)) * jnp.sqrt(cos_term + sin_term)
+
+    periodic_r_ee = periodic_r_ee * (1.0 - jnp.eye(n))
+    return periodic_r_ee[..., None]
+
+  return apply
+
+
 def _jastrow_ee(
     r_ee: jnp.ndarray,
     params: ParamTree,
     nspins: tuple[int, int],
+    ndim: int,
     jastrow_fun: Callable[[jnp.ndarray, float, jnp.ndarray], jnp.ndarray],
 ) -> jnp.ndarray:
   """Jastrow factor for electron-electron cusps."""
@@ -47,22 +79,29 @@ def _jastrow_ee(
 
   if r_ees_parallel.shape[0] > 0:
     jastrow_ee_par = jnp.sum(
-        jastrow_fun(r_ees_parallel, 0.25, params['ee_par'])
+        jastrow_fun(r_ees_parallel, 1 / (ndim + 1), params['ee_par'])
     )
   else:
     jastrow_ee_par = jnp.asarray(0.0)
 
   if r_ees[0][1].shape[0] > 0:
-    jastrow_ee_anti = jnp.sum(jastrow_fun(r_ees[0][1], 0.5, params['ee_anti']))
+    jastrow_ee_anti = jnp.sum(
+        jastrow_fun(r_ees[0][1], 1 / (ndim - 1), params['ee_anti']))
   else:
     jastrow_ee_anti = jnp.asarray(0.0)
 
   return jastrow_ee_anti + jastrow_ee_par
 
 
-def make_simple_ee_jastrow():
+def make_simple_ee_jastrow(lattice: Optional[jnp.ndarray] = None, ndim: int = 3):
   """Creates a Jastrow factor for electron-electron cusps."""
 
+  # If working in PBC, use periodic distance for the Jastrow
+  if lattice is not None:
+    norm = make_periodic_r_ee(lattice)
+  else:
+    norm = lambda x: jnp.linalg.norm(x, axis = -1, keepdims = True)
+
   def simple_ee_cusp_fun(
       r: jnp.ndarray, cusp: float, alpha: jnp.ndarray
   ) -> jnp.ndarray:
@@ -80,20 +119,24 @@ def init() -> Mapping[str, jnp.ndarray]:
     return params
 
   def apply(
-      r_ee: jnp.ndarray,
+      ee: jnp.ndarray,
       params: ParamTree,
       nspins: tuple[int, int],
   ) -> jnp.ndarray:
     """Jastrow factor for electron-electron cusps."""
-    return _jastrow_ee(r_ee, params, nspins, jastrow_fun=simple_ee_cusp_fun)
+    r_ee = norm(ee)
+    return _jastrow_ee(r_ee, params, nspins, ndim, jastrow_fun=simple_ee_cusp_fun)
 
   return init, apply
 
 
-def get_jastrow(jastrow: JastrowType):
+def get_jastrow(
+    jastrow: JastrowType, 
+    lattice: Optional[jnp.ndarray] = None, 
+    ndim: int = 3):
   jastrow_init, jastrow_apply = None, None
   if jastrow == JastrowType.SIMPLE_EE:
-    jastrow_init, jastrow_apply = make_simple_ee_jastrow()
+    jastrow_init, jastrow_apply = make_simple_ee_jastrow(lattice, ndim)
   elif jastrow != JastrowType.NONE:
     raise ValueError(f'Unknown Jastrow Factor type: {jastrow}')
 

ferminet/mcmc.py

@@ -25,6 +25,20 @@
 from jax import lax
 from jax import numpy as jnp
 import numpy as np
+from typing import Callable
+
+
+def map_to_simulation_cell(pos, lattice, rec, ndim):
+  """If working in PBC, map electrons back to the supercell
+  """
+  pos_ = jnp.reshape(pos, [-1, ndim])
+  phase = jnp.einsum('il,kl->ki', rec / (2 * jnp.pi), pos_)
+  phase_prim = phase % 1
+  prim = jnp.einsum('il,kl->ki', lattice, phase_prim)
+  return prim.flatten()
+
+batch_map_to_simulation_cell = jax.vmap(
+    map_to_simulation_cell, in_axes = (0, None, None, None))
 
 
 def _harmonic_mean(x, atoms):
@@ -78,6 +92,7 @@ def mh_accept(x1, x2, lp_1, lp_2, ratio, key, num_accepts):
 def mh_update(
     params: networks.ParamTree,
     f: networks.LogFermiNetLike,
+    _map: Callable,
     data: networks.FermiNetData,
     key: chex.PRNGKey,
     lp_1,
@@ -94,6 +109,8 @@ def mh_update(
     params: Wavefuncttion parameters.
     f: Callable with signature f(params, x) which returns the log of the
       wavefunction (i.e. the sqaure root of the log probability of x).
+    _map: If PBC, callable which maps electrons back to the simulation cell.
+      If OBC, identity function.
     data: Initial MCMC configurations (batched).
     key: RNG state.
     lp_1: log probability of f evaluated at x1 given parameters params.
@@ -120,6 +137,7 @@ def mh_update(
   x1 = data.positions
   if atoms is None:  # symmetric proposal, same stddev everywhere
     x2 = x1 + stddev * jax.random.normal(subkey, shape=x1.shape)  # proposal
+    x2 = _map(x2)
     lp_2 = 2.0 * f(
         params, x2, data.spins, data.atoms, data.charges
     )  # log prob of proposal
@@ -130,6 +148,7 @@ def mh_update(
     hmean1 = _harmonic_mean(x1, atoms)  # harmonic mean of distances to nuclei
 
     x2 = x1 + stddev * hmean1 * jax.random.normal(subkey, shape=x1.shape)
+    x2 = _map(x2)
     lp_2 = 2.0 * f(
         params, x2, data.spins, data.atoms, data.charges
     )  # log prob of proposal
@@ -150,6 +169,7 @@ def mh_update(
 def mh_block_update(
     params: networks.ParamTree,
     f: networks.LogFermiNetLike,
+    _map: Callable,
     data: networks.FermiNetData,
     key: chex.PRNGKey,
     lp_1,
@@ -166,6 +186,8 @@ def mh_block_update(
     params: Wavefuncttion parameters.
     f: Callable with LogFermiNetLike signature which returns the log of the
       wavefunction (i.e. the sqaure root of the log probability of x).
+    _map: If PBC, callable which maps electrons back to the simulation cell.
+      If OBC, identity function.
     data: Initial MCMC configuration (batched).
     key: RNG state.
     lp_1: log probability of f evaluated at x1 given parameters params.
@@ -199,6 +221,7 @@ def mh_block_update(
     x2 = x1.at[:, ii].add(
         stddev * jax.random.normal(subkey, shape=x1[:, ii].shape))
     x2 = jnp.reshape(x2, [batch_size, -1])
+    x2 = _map(x2)
     if pad > 0:
       x2 = x2[..., :-pad*ndim]
     # log prob of proposal
@@ -222,7 +245,8 @@ def make_mcmc_step(batch_network,
                    steps=10,
                    atoms=None,
                    ndim=3,
-                   blocks=1):
+                   blocks=1,
+                   lattice=None):
   """Creates the MCMC step function.
 
   Args:
@@ -238,12 +262,22 @@ def make_mcmc_step(batch_network,
     ndim: Dimensionality of the system (usually 3).
     blocks: Number of blocks to split the updates into. If 1, use all-electron
       moves.
+    lattice: If None, assume OBC. Otherwise matrix with supercell lattice 
+      vectors
 
   Returns:
     Callable which performs the set of MCMC steps.
   """
   inner_fun = mh_block_update if blocks > 1 else mh_update
 
+  # If PBC, map electrons back to the supercell
+  if lattice is not None:
+    rec = 2 * jnp.pi * jnp.linalg.inv(lattice)
+    _map = lambda p: batch_map_to_simulation_cell(p, lattice, rec, ndim)
+  else:
+    rec = None
+    _map = lambda p: p
+
   def mcmc_step(params, data, key, width):
     """Performs a set of MCMC steps.
 
@@ -263,6 +297,7 @@ def step_fn(i, x):
       return inner_fun(
           params,
           batch_network,
+          _map,
           *x,
           stddev=width,
           atoms=atoms,

ferminet/networks.py

@@ -267,6 +267,8 @@ class BaseNetworkOptions:
   Attributes:
     ndim: dimension of system. Change only with caution.
     determinants: Number of determinants to use.
+    lattice: If None, assume OBC. Otherwise matrix with supercell lattice 
+      vectors.
     states: Number of outputs, one per excited (or ground) state. Ignored if 0.
     full_det: If true, evaluate determinants over all electrons. Otherwise,
       block-diagonalise determinants into spin channels.
@@ -281,6 +283,7 @@ class BaseNetworkOptions:
   """
 
   ndim: int = 3
+  lattice: Optional[jnp.ndarray] = None
   determinants: int = 16
   states: int = 0
   full_det: bool = True
@@ -1079,7 +1082,8 @@ def make_orbitals(
   equivariant_layers_init, equivariant_layers_apply = equivariant_layers
 
   # Optional Jastrow factor.
-  jastrow_init, jastrow_apply = jastrows.get_jastrow(options.jastrow)
+  jastrow_init, jastrow_apply = jastrows.get_jastrow(
+      options.jastrow, options.lattice, options.ndim)
 
   def init(key: chex.PRNGKey) -> ParamTree:
     """Returns initial random parameters for creating orbitals.
@@ -1235,7 +1239,7 @@ def apply(
     # Added pre-determinant for compatibility with pretraining.
     if jastrow_apply is not None:
       jastrow = jnp.exp(
-          jastrow_apply(r_ee, params['jastrow'], nspins) / sum(nspins)
+          jastrow_apply(ee, params['jastrow'], nspins) / sum(nspins)
       )
       orbitals = [orbital * jastrow for orbital in orbitals]
 
@@ -1365,6 +1369,7 @@ def make_fermi_net(
     charges: jnp.ndarray,
     *,
     ndim: int = 3,
+    lattice: Optional[jnp.ndarray] = None,
     determinants: int = 16,
     states: int = 0,
     envelope: Optional[envelopes.Envelope] = None,
@@ -1389,6 +1394,8 @@ def make_fermi_net(
     nspins: Tuple of the number of spin-up and spin-down electrons.
     charges: (natom) array of atom nuclear charges.
     ndim: dimension of system. Change only with caution.
+    lattice: If None, assume OBC. Otherwise matrix with supercell lattice 
+      vectors.
     determinants: Number of determinants to use.
     states: Number of outputs, one per excited (or ground) state. Ignored if 0.
     envelope: Envelope to use to impose orbitals go to zero at infinity.
@@ -1448,6 +1455,7 @@ def make_fermi_net(
 
   options = FermiNetOptions(
       ndim=ndim,
+      lattice=lattice,
       determinants=determinants,
       states=states,
       rescale_inputs=rescale_inputs,