diff --git a/ferminet/base_config.py b/ferminet/base_config.py index e491143..0c8a32f 100644 --- a/ferminet/base_config.py +++ b/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. diff --git a/ferminet/envelopes.py b/ferminet/envelopes.py index b93adc2..860d187 100644 --- a/ferminet/envelopes.py +++ b/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) diff --git a/ferminet/hamiltonian.py b/ferminet/hamiltonian.py index d10ae2c..cb639af 100644 --- a/ferminet/hamiltonian.py +++ b/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) diff --git a/ferminet/pretrain.py b/ferminet/pretrain.py index 7a35dff..970eb79 100644 --- a/ferminet/pretrain.py +++ b/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) diff --git a/ferminet/train.py b/ferminet/train.py index 41333af..6fa4835 100644 --- a/ferminet/train.py +++ b/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: