feat: Add support for fermions in a harmonic oscillator potential

ferminet/base_config.py

@@ -144,6 +144,10 @@ def default() -> ml_collections.ConfigDict:
       },
       'system': {
           'type': SystemType.MOLECULE.value,
+          'potential': 'coulomb',  # One of 'coulomb' or 'harm'
+          'omega': 1.0,  # Harmonic oscillator frequency
+          'interacting': True,  # Whether electrons interact
+          'scale_alpha': 1.0,  # Scale factor for electron-electron interaction
           # Specify the system.
           # 1. Specify the system by setting variables below.
           # list of system.Atom objects with element type and position.

ferminet/envelopes.py

@@ -41,6 +41,7 @@ class EnvelopeLabel(enum.Enum):
   NULL = enum.auto()
   STO = enum.auto()
   STO_POLY = enum.auto()
+  GAUSSIAN = enum.auto()
 
 
 class EnvelopeInit(Protocol):
@@ -124,6 +125,36 @@ def apply(*, ae: jnp.ndarray, r_ae: jnp.ndarray, r_ee: jnp.ndarray,
   return Envelope(EnvelopeType.PRE_DETERMINANT, init, apply)
 
 
+def make_gaussian_envelope(**kwargs) -> Envelope:
+  """Creates an isotropic Gaussian exponentially decaying multiplicative envelope."""
+  omega = kwargs.get('omega', 1.0)
+
+  def init(
+      natom: int, output_dims: Sequence[int], ndim: int = 3
+  ) -> Sequence[Mapping[str, jnp.ndarray]]:
+    del ndim  # unused
+    params = []
+    # Ground state gaussian orbital: (omega/pi)^0.75 * exp(-0.5 * omega * r^2)
+    # Our envelope is pi * exp(-(r * sigma)^2)
+    # So sigma = sqrt(omega / 2) and pi = (omega/pi)^0.75
+    pi_init = (omega / jnp.pi)**0.75
+    sigma_init = jnp.sqrt(omega / 2.0)
+    for output_dim in output_dims:
+      params.append({
+          'pi': jnp.ones(shape=(natom, output_dim)) * pi_init,
+          'sigma': jnp.ones(shape=(natom, output_dim)) * sigma_init
+      })
+    return params
+
+  def apply(*, ae: jnp.ndarray, r_ae: jnp.ndarray, r_ee: jnp.ndarray,
+            pi: jnp.ndarray, sigma: jnp.ndarray) -> jnp.ndarray:
+    """Computes an isotropic Gaussian exponentially-decaying multiplicative envelope."""
+    del ae, r_ee  # unused
+    return jnp.sum(jnp.exp(-(r_ae * sigma)**2) * pi, axis=1)
+
+  return Envelope(EnvelopeType.PRE_DETERMINANT, init, apply)
+
+
 def make_bottleneck_envelope(nenv: int = 16) -> Envelope:
   """Each orbital has a linear projection of a small number of envelopes.
 
@@ -314,5 +345,6 @@ def get_envelope(
       EnvelopeLabel.DIAGONAL: make_diagonal_envelope,
       EnvelopeLabel.FULL: make_full_envelope,
       EnvelopeLabel.NULL: make_null_envelope,
+      EnvelopeLabel.GAUSSIAN: make_gaussian_envelope,
   }
   return envelope_builders[envelope_label](**kwargs)

ferminet/hamiltonian.py

@@ -299,21 +299,32 @@ def potential_nuclear_nuclear(charges: Array, atoms: Array) -> jnp.ndarray:
 
 
 def potential_energy(r_ae: Array, r_ee: Array, atoms: Array,
-                     charges: Array) -> jnp.ndarray:
+                     charges: Array, potential_type: str = 'coulomb',
+                     omega: float = 0.0, scale_alpha: float = 1.0) -> jnp.ndarray:
   """Returns the potential energy for this electron configuration.
 
   Args:
-    r_ae: Shape (nelectrons, natoms). r_ae[i, j] gives the distance between
+    r_ae: Shape (nelectrons, natoms, ...). r_ae[i, j, 0] gives the distance between
       electron i and atom j.
     r_ee: Shape (neletrons, nelectrons, :). r_ee[i,j,0] gives the distance
       between electrons i and j. Other elements in the final axes are not
       required.
     atoms: Shape (natoms, ndim). Positions of the atoms.
     charges: Shape (natoms). Nuclear charges of the atoms.
+    potential_type: String specifying potential ('coulomb' or 'harm').
+    omega: Harmonic oscillator frequency.
+    scale_alpha: Scaling factor for electron-electron interaction.
   """
-  return (potential_electron_electron(r_ee) +
-          potential_electron_nuclear(charges, r_ae) +
-          potential_nuclear_nuclear(charges, atoms))
+  v_ee = potential_electron_electron(r_ee) * scale_alpha
+  if potential_type == 'harm':
+    # Harmonic oscillator potential: 1/2 \omega^2 \sum_i |r_i - R|^2
+    v_ae = 0.5 * (omega**2) * jnp.sum(r_ae[:, 0, 0]**2)
+    v_aa = 0.0
+  else:
+    v_ae = potential_electron_nuclear(charges, r_ae)
+    v_aa = potential_nuclear_nuclear(charges, atoms)
+
+  return v_ee + v_ae + v_aa
 
 
 def local_energy(
@@ -328,6 +339,9 @@ def local_energy(
     state_specific: bool = False,
     pp_type: str = 'ccecp',
     pp_symbols: Sequence[str] | None = None,
+    potential_type: str = 'coulomb',
+    omega: float = 0.0,
+    scale_alpha: float = 1.0,
 ) -> LocalEnergy:
   """Creates the function to evaluate the local energy.
 
@@ -351,6 +365,9 @@ def local_energy(
       provided.
     pp_symbols: sequence of element symbols for which the pseudopotential is
       used.
+    potential_type: String specifying potential ('coulomb' or 'harm').
+    omega: Harmonic oscillator frequency.
+    scale_alpha: Scaling factor for electron-electron interaction.
 
   Returns:
     Callable with signature e_l(params, key, data) which evaluates the local
@@ -397,9 +414,9 @@ def _e_l(
       ae, _, r_ae, r_ee = vmap_features(positions, data.atoms, ndim)
 
       # Compute potential energy
-      vmap_pot = jax.vmap(potential_energy, (0, 0, None, None))
+      vmap_pot = jax.vmap(potential_energy, (0, 0, None, None, None, None, None))
       pot_spectrum = vmap_pot(
-          r_ae, r_ee, data.atoms, effective_charges)[:, None]
+          r_ae, r_ee, data.atoms, effective_charges, potential_type, omega, scale_alpha)[:, None]
 
       if use_pp:
         data_vmap_dims = networks.FermiNetData(
@@ -452,7 +469,7 @@ def _e_l(
           data.atoms,
           ndim,
       )
-      potential = (potential_energy(r_ae, r_ee, data.atoms, effective_charges) +
+      potential = (potential_energy(r_ae, r_ee, data.atoms, effective_charges, potential_type, omega, scale_alpha) +
                    pp_local(r_ae) +
                    pp_nonlocal(key, f, params, data, ae, r_ae))
       kinetic = ke(params, data)

ferminet/pretrain.py

@@ -31,6 +31,81 @@
 import pyscf
 
 
+import math
+
+class HarmonicOscillatorExact:
+  """Exact orbitals for a 3D isotropic harmonic oscillator."""
+
+  def __init__(self, omega: float):
+    self.omega = omega
+
+  def eval_orbitals(self, pos: Union[np.ndarray, jnp.ndarray], nspins: Tuple[int, int]) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Evaluates harmonic oscillator orbitals at a set of positions."""
+    if not isinstance(pos, jnp.ndarray) and not isinstance(pos, np.ndarray):
+      pos = jnp.asarray(pos)
+
+    leading_dims = pos.shape[:-1]
+    pos = jnp.reshape(pos, leading_dims + (sum(nspins), 3))
+
+    n_orbs = max(nspins)
+
+    orbitals = []
+    for n in range(10):
+      for nx in range(n + 1):
+        for ny in range(n - nx + 1):
+          nz = n - nx - ny
+          orbitals.append((nx, ny, nz))
+          if len(orbitals) >= n_orbs:
+            break
+        if len(orbitals) >= n_orbs:
+            break
+      if len(orbitals) >= n_orbs:
+          break
+
+    def hermite(n, x):
+      if n == 0: return jnp.ones_like(x)
+      if n == 1: return 2.0 * x
+      if n == 2: return 4.0 * x**2 - 2.0
+      if n == 3: return 8.0 * x**3 - 12.0 * x
+      if n == 4: return 16.0 * x**4 - 48.0 * x**2 + 12.0
+      if n == 5: return 32.0 * x**5 - 160.0 * x**3 + 120.0 * x
+      raise NotImplementedError(f"Hermite polynomial of order {n} not implemented.")
+
+    def eval_orb(nx, ny, nz, r):
+        x, y, z = r[..., 0], r[..., 1], r[..., 2]
+        sq_omega = jnp.sqrt(self.omega)
+        hx = hermite(nx, sq_omega * x)
+        hy = hermite(ny, sq_omega * y)
+        hz = hermite(nz, sq_omega * z)
+
+        norm_x = 1.0 / jnp.sqrt(2**nx * math.factorial(nx))
+        norm_y = 1.0 / jnp.sqrt(2**ny * math.factorial(ny))
+        norm_z = 1.0 / jnp.sqrt(2**nz * math.factorial(nz))
+
+        gaussian = jnp.exp(-0.5 * self.omega * (x**2 + y**2 + z**2))
+        return norm_x * norm_y * norm_z * (self.omega / jnp.pi)**0.75 * hx * hy * hz * gaussian
+
+    evals = [eval_orb(nx, ny, nz, pos) for nx, ny, nz in orbitals]
+    evals = jnp.stack(evals, axis=-1)
+
+    alpha_spin = evals[..., :nspins[0], :nspins[0]]
+    beta_spin = evals[..., nspins[0]:, :nspins[1]]
+
+    return alpha_spin, beta_spin
+
+  def eval_slater(self,
+                  pos: Union[jnp.ndarray, np.ndarray],
+                  nspins: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
+    """Evaluates the Slater determinant."""
+    matrices = self.eval_orbitals(pos, nspins)
+    slogdets = [jnp.linalg.slogdet(elem) for elem in matrices]
+    sign_alpha, sign_beta = [elem[0] for elem in slogdets]
+    log_abs_wf_alpha, log_abs_wf_beta = [elem[1] for elem in slogdets]
+    log_abs_slater_determinant = log_abs_wf_alpha + log_abs_wf_beta
+    sign = sign_alpha * sign_beta
+    return sign, log_abs_slater_determinant
+
+
 def get_hf(molecule: Sequence[system.Atom] | None = None,
            nspins: Tuple[int, int] | None = None,
            basis: str | None = 'sto-3g',
@@ -39,7 +114,9 @@ def get_hf(molecule: Sequence[system.Atom] | None = None,
            pyscf_mol: pyscf.gto.Mole | None = None,
            restricted: bool | None = False,
            states: int = 0,
-           excitation_type: str = 'ordered') -> scf.Scf:
+           excitation_type: str = 'ordered',
+           potential_type: str = 'coulomb',
+           omega: float = 1.0) -> Union[scf.Scf, HarmonicOscillatorExact]:
   """Returns an Scf object with the Hartree-Fock solution to the system.
 
   Args:
@@ -59,7 +136,12 @@ def get_hf(molecule: Sequence[system.Atom] | None = None,
     excitation_type: The way to construct different states for excited state
       pretraining. One of 'ordered' or 'random'. 'Ordered' tends to work better,
       but 'random' is necessary for some systems, especially double excitaitons.
+    potential_type: The type of potential ('coulomb' or 'harm').
+    omega: The harmonic oscillator frequency (only used if potential_type is 'harm').
   """
+  if potential_type == 'harm':
+    return HarmonicOscillatorExact(omega)
+
   if pyscf_mol:
     scf_approx = scf.Scf(pyscf_mol=pyscf_mol,
                          restricted=restricted)

ferminet/train.py

@@ -468,11 +468,14 @@ def train(cfg: ml_collections.ConfigDict, writer_manager=None):
         ecp=ecp,
         core_electrons=core_electrons,
         states=cfg.system.states,
-        excitation_type=cfg.pretrain.get('excitation_type', 'ordered'))
+        excitation_type=cfg.pretrain.get('excitation_type', 'ordered'),
+        potential_type=cfg.system.get('potential', 'coulomb'),
+        omega=cfg.system.get('omega', 1.0))
     # broadcast the result of PySCF from host 0 to all other hosts
-    hartree_fock.mean_field.mo_coeff = multihost_utils.broadcast_one_to_all(
-        hartree_fock.mean_field.mo_coeff
-    )
+    if cfg.system.get('potential', 'coulomb') != 'harm':
+      hartree_fock.mean_field.mo_coeff = multihost_utils.broadcast_one_to_all(
+          hartree_fock.mean_field.mo_coeff
+      )
 
   if cfg.network.make_feature_layer_fn:
     feature_layer_module, feature_layer_fn = (
@@ -501,7 +504,10 @@ def train(cfg: ml_collections.ConfigDict, writer_manager=None):
     make_envelope = getattr(envelope_module, envelope_fn)
     envelope = make_envelope(**cfg.network.make_envelope_kwargs)  # type: envelopes.Envelope
   else:
-    envelope = envelopes.make_isotropic_envelope()
+    if cfg.system.get('potential', 'coulomb') == 'harm':
+      envelope = envelopes.make_gaussian_envelope(omega=cfg.system.get('omega', 1.0))
+    else:
+      envelope = envelopes.make_isotropic_envelope()
 
   use_complex = cfg.network.get('complex', False)
   if cfg.network.network_type == 'ferminet':
@@ -615,13 +621,16 @@ def log_network(*args, **kwargs):
     key, subkey = jax.random.split(key)
     # make sure data on each host is initialized differently
     subkey = jax.random.fold_in(subkey, jax.process_index())
+    init_width = cfg.mcmc.init_width
+    if cfg.system.get('potential', 'coulomb') == 'harm':
+      init_width /= np.sqrt(cfg.system.get('omega', 1.0))
     # create electron state (position and spin)
     pos, spins = init_electrons(
         subkey,
         cfg.system.molecule,
         cfg.system.electrons,
         batch_size=total_host_batch_size,
-        init_width=cfg.mcmc.init_width,
+        init_width=init_width,
         core_electrons=core_electrons,
     )
     # For excited states, each device has a batch of walkers, where each walker
@@ -764,6 +773,9 @@ def log_network(*args, **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,
+      potential_type=cfg.system.get('potential', 'coulomb'),
+      omega=cfg.system.get('omega', 1.0),
+      scale_alpha=cfg.system.get('scale_alpha', 1.0) if cfg.system.get('interacting', True) else 0.0,
       **cfg.system.make_local_energy_kwargs,
   )
 
@@ -907,8 +919,11 @@ def learning_rate_schedule(t_: jnp.ndarray) -> jnp.ndarray:
   if mcmc_width_ckpt is not None:
     mcmc_width = kfac_jax.utils.replicate_all_local_devices(mcmc_width_ckpt[0])
   else:
+    move_width = cfg.mcmc.move_width
+    if cfg.system.get('potential', 'coulomb') == 'harm':
+      move_width /= np.sqrt(cfg.system.get('omega', 1.0))
     mcmc_width = kfac_jax.utils.replicate_all_local_devices(
-        jnp.asarray(cfg.mcmc.move_width))
+        jnp.asarray(move_width))
   pmoves = np.zeros(cfg.mcmc.adapt_frequency)
 
   if t_init == 0: