Remove NotImplementedError and add missing arguments to custom Hamiltonians

Author: APerezFadon

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,
@@ -188,6 +191,9 @@ def local_energy(
   """
   if states:
     raise NotImplementedError('Excited states not implemented with PBC.')
+  if pp_symbols:
+    raise NotImplementedError('Pseudopotentials not implemented with PBC.')
+  
   del nspins
   if lattice is None:
     lattice = jnp.eye(3)

ferminet/train.py

@@ -743,12 +743,6 @@ def log_network(*args, **kwargs):
   # Construct loss and optimizer
   laplacian_method = cfg.optim.get('laplacian', 'default')
   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,7 +752,12 @@ 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),
+        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')