Merge pull request #110 from APerezFadon:main

PiperOrigin-RevId: 914789354
Change-Id: Ie5b25ec5cd6e2f8c7be1e588eb832699cdb3f374

Author: jsspencer

ferminet/pbc/hamiltonian.py

@@ -158,6 +158,9 @@ def local_energy(
     complex_output: bool = False,
     laplacian_method: str = 'default',
     states: int = 0,
+    state_specific: bool = False,
+    pp_type: str = 'ccecp',
+    pp_symbols: Sequence[str] | None = None,
     lattice: Optional[jnp.ndarray] = None,
     heg: bool = True,
     convergence_radius: int = 5,
@@ -176,6 +179,10 @@ def local_energy(
       'folx': use Microsoft's implementation of forward laplacian
     states: Number of excited states to compute. Not implemented, only present
       for consistency of calling convention.
+    state_specific: Not implemented.
+    pp_type: type of pseudopotential to use. Not implemented.
+    pp_symbols: sequence of element symbols for which the pseudopotential is
+      used. Not implemented.
     lattice: Shape (ndim, ndim). Matrix of lattice vectors. Default: identity
       matrix.
     heg: bool. Flag to enable features specific to the electron gas.
@@ -186,9 +193,13 @@ def local_energy(
     energy of the wavefunction given the parameters params, RNG state key,
     and a single MCMC configuration in data.
   """
-  if states:
+  if states > 0 or state_specific:
     raise NotImplementedError('Excited states not implemented with PBC.')
+  if pp_symbols:
+    raise NotImplementedError('Pseudopotentials not implemented with PBC.')
+
   del nspins
+  del pp_type
   if lattice is None:
     lattice = jnp.eye(3)
 

ferminet/train.py

@@ -742,13 +742,8 @@ def log_network(*args, **kwargs):
   )
   # 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:
-    if laplacian_method != 'default':
-      raise NotImplementedError(f'Laplacian method {laplacian_method}'
-                                'not yet supported by custom local energy fns.')
-    if cfg.optim.objective == 'vmc_overlap':
-      raise NotImplementedError('Overlap penalty not yet supported for custom'
-                                'local energy fns.')
     local_energy_module, local_energy_fn = (
         cfg.system.make_local_energy_fn.rsplit('.', maxsplit=1))
     local_energy_module = importlib.import_module(local_energy_module)
@@ -758,10 +753,15 @@ def log_network(*args, **kwargs):
         charges=charges,
         nspins=nspins,
         use_scan=False,
+        complex_output=use_complex,
+        laplacian_method=laplacian_method,
         states=cfg.system.get('states', 0),
-        **cfg.system.make_local_energy_kwargs)
+        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:
-    pp_symbols = cfg.system.get('pp', {'symbols': None}).get('symbols')
     local_energy_fn = hamiltonian.local_energy(
         f=signed_network,
         charges=charges,