Fix for general ndim in OBC

Author: APerezFadon

ferminet/hamiltonian.py

@@ -321,6 +321,7 @@ def local_energy(
     charges: jnp.ndarray,
     nspins: Sequence[int],
     use_scan: bool = False,
+    ndim: int = 3,
     complex_output: bool = False,
     laplacian_method: str = 'default',
     states: int = 0,
@@ -336,6 +337,7 @@ def local_energy(
     charges: Shape (natoms). Nuclear charges of the atoms.
     nspins: Number of particles of each spin.
     use_scan: Whether to use a `lax.scan` for computing the laplacian.
+    ndim: Number of dimensions.
     complex_output: If true, the output of f is complex-valued.
     laplacian_method: Laplacian calculation method. One of:
       'default': take jvp(grad), looping over inputs
@@ -360,7 +362,7 @@ def local_energy(
   if not pp_symbols:
     effective_charges = charges
     use_pp = False
-  else:
+  elif ndim == 3:
     effective_charges, pp_local, pp_nonlocal = pp.make_pp_potential(
         charges=charges,
         symbols=pp_symbols,
@@ -369,6 +371,8 @@ def local_energy(
         complex_output=complex_output
     )
     use_pp = not jnp.all(effective_charges == charges)
+  else:
+    raise NotImplementedError("Pseudopotentials only implemented for 3D")
 
   if not use_pp:
     pp_local = lambda *args, **kwargs: 0.0
@@ -386,9 +390,10 @@ def _e_l(
     """
     if states:
       # Compute features
-      vmap_features = jax.vmap(networks.construct_input_features, (0, None))
+      vmap_features = jax.vmap(
+          networks.construct_input_features, (0, None, None))
       positions = jnp.reshape(data.positions, [states, -1])
-      ae, _, r_ae, r_ee = vmap_features(positions, data.atoms)
+      ae, _, r_ae, r_ee = vmap_features(positions, data.atoms, ndim)
 
       # Compute potential energy
       vmap_pot = jax.vmap(potential_energy, (0, 0, None, None))
@@ -442,7 +447,7 @@ def _e_l(
                                 complex_output=complex_output,
                                 laplacian_method=laplacian_method)
       ae, _, r_ae, r_ee = networks.construct_input_features(
-          data.positions, data.atoms
+          data.positions, data.atoms, ndim,
       )
       potential = (potential_energy(r_ae, r_ee, data.atoms, effective_charges) +
                    pp_local(r_ae) +

ferminet/train.py

@@ -739,40 +739,32 @@ def log_network(*args, **kwargs):
       steps=cfg.mcmc.steps,
       atoms=atoms_to_mcmc,
       blocks=cfg.mcmc.blocks * num_states,
+      ndim=cfg.system.ndim,
   )
+  
   # Construct loss and optimizer
-  laplacian_method = cfg.optim.get('laplacian', 'default')
-  pp_symbols = cfg.system.get('pp', {'symbols': None}).get('symbols')
   if cfg.system.make_local_energy_fn:
     local_energy_module, local_energy_fn = (
         cfg.system.make_local_energy_fn.rsplit('.', maxsplit=1))
     local_energy_module = importlib.import_module(local_energy_module)
     make_local_energy = getattr(local_energy_module, local_energy_fn)  # type: hamiltonian.MakeLocalEnergy
-    local_energy_fn = make_local_energy(
-        f=signed_network,
-        charges=charges,
-        nspins=nspins,
-        use_scan=False,
-        complex_output=use_complex,
-        laplacian_method=laplacian_method,
-        states=cfg.system.get('states', 0),
-        state_specific=(cfg.optim.objective == 'vmc_overlap'),
-        pp_type=cfg.system.get('pp', {'type': 'ccecp'}).get('type'),
-        pp_symbols=pp_symbols if cfg.system.get('use_pp') else None,
-        **cfg.system.make_local_energy_kwargs,
-    )
   else:
-    local_energy_fn = hamiltonian.local_energy(
-        f=signed_network,
-        charges=charges,
-        nspins=nspins,
-        use_scan=False,
-        complex_output=use_complex,
-        laplacian_method=laplacian_method,
-        states=cfg.system.get('states', 0),
-        state_specific=(cfg.optim.objective == 'vmc_overlap'),
-        pp_type=cfg.system.get('pp', {'type': 'ccecp'}).get('type'),
-        pp_symbols=pp_symbols if cfg.system.get('use_pp') else None)
+    make_local_energy = hamiltonian.local_energy
+  laplacian_method = cfg.optim.get('laplacian', 'default')
+  pp_symbols = cfg.system.get('pp', {'symbols': None}).get('symbols')
+  local_energy_fn = make_local_energy(
+      f=signed_network,
+      charges=charges,
+      nspins=nspins,
+      use_scan=False,
+      ndim=cfg.system.ndim,
+      complex_output=use_complex,
+      laplacian_method=laplacian_method,
+      states=cfg.system.get('states', 0),
+      state_specific=(cfg.optim.objective == 'vmc_overlap'),
+      pp_type=cfg.system.get('pp', {'type': 'ccecp'}).get('type'),
+      pp_symbols=pp_symbols if cfg.system.get('use_pp') else None,
+      **cfg.system.make_local_energy_kwargs)
 
   if cfg.optim.get('spin_energy', 0.0) > 0.0:
     # Minimize <H + c * S^2> instead of just <H>