From 985c1d8c0a5657a7e2ecb8a1f3905718e0e3591f Mon Sep 17 00:00:00 2001 From: Andres Date: Mon, 25 May 2026 15:53:05 +0100 Subject: [PATCH] Integrate PBC to the config --- ferminet/base_config.py | 6 +- ferminet/configs/heg.py | 44 +++++----- ferminet/jastrows.py | 59 ++++++++++++-- ferminet/mcmc.py | 37 ++++++++- ferminet/networks.py | 12 ++- ferminet/pbc/envelopes.py | 9 ++- ferminet/pbc/ewald2d.py | 146 +++++++++++++++++++++++++++++++++ ferminet/pbc/ewald3d.py | 147 ++++++++++++++++++++++++++++++++++ ferminet/pbc/feature_layer.py | 3 +- ferminet/pbc/hamiltonian.py | 144 ++++----------------------------- ferminet/psiformer.py | 4 + ferminet/train.py | 35 +++++--- 12 files changed, 466 insertions(+), 180 deletions(-) create mode 100644 ferminet/pbc/ewald2d.py create mode 100644 ferminet/pbc/ewald3d.py diff --git a/ferminet/base_config.py b/ferminet/base_config.py index e491143..5f883b9 100644 --- a/ferminet/base_config.py +++ b/ferminet/base_config.py @@ -148,10 +148,12 @@ def default() -> ml_collections.ConfigDict: # 1. Specify the system by setting variables below. # list of system.Atom objects with element type and position. 'molecule': config_dict.placeholder(list), + # If None, assume OBC, if PBC then matrix with the + # supercell lattice vectors + 'lattice': None, # number of spin up, spin-down electrons 'electrons': tuple(), - # Dimensionality. Change with care. FermiNet implementation currently - # assumes 3D systems. + # Dimensionality. 'ndim': 3, # Number of excited states. If 0, use normal ground state machinery. # If 1, compute ground state using excited state machinery. If >1, diff --git a/ferminet/configs/heg.py b/ferminet/configs/heg.py index 3590aae..96ab7a3 100644 --- a/ferminet/configs/heg.py +++ b/ferminet/configs/heg.py @@ -15,43 +15,41 @@ """Unpolarised 14 electron simple cubic homogeneous electron gas.""" from ferminet import base_config -from ferminet.pbc import envelopes from ferminet.utils import system import numpy as np -def _sc_lattice_vecs(rs: float, nelec: int) -> np.ndarray: +def _sc_lattice_vecs(rs: float, nelec: int, ndim: int) -> np.ndarray: """Returns simple cubic lattice vectors with Wigner-Seitz radius rs.""" - volume = (4 / 3) * np.pi * (rs**3) * nelec - length = volume**(1 / 3) - return length * np.eye(3) - + if ndim == 2: + area = np.pi * (rs**2) * nelec + length = area**(1 / 2) + return length * np.eye(2) + elif ndim == 3: + area = 4 * np.pi * (rs**3) * nelec / 3 + length = area**(1 / 3) + return length * np.eye(3) + else: + raise NotImplementedError + +rs = 1.0 def get_config(): - """Returns config for running unpolarised 14 electron gas with FermiNet.""" # Get default options. cfg = base_config.default() + + # SYSTEM cfg.system.electrons = (7, 7) # A ghost atom at the origin defines one-electron coordinate system. # Element 'X' is a dummy nucleus with zero charge - cfg.system.molecule = [system.Atom("X", (0., 0., 0.))] + cfg.system.molecule = [system.Atom('X', (0, 0, 0))] + + cfg.system.lattice = _sc_lattice_vecs(rs, sum(cfg.system.electrons), cfg.system.ndim) + cfg.system.make_local_energy_kwargs['heg'] = True + cfg.network.make_feature_layer_kwargs['include_r_ae'] = False + # Pretraining is not currently implemented for systems in PBC cfg.pretrain.method = None - lattice = _sc_lattice_vecs(1.0, sum(cfg.system.electrons)) - kpoints = envelopes.make_kpoints(lattice, cfg.system.electrons) - - cfg.system.make_local_energy_fn = "ferminet.pbc.hamiltonian.local_energy" - cfg.system.make_local_energy_kwargs = {"lattice": lattice, "heg": True} - cfg.network.make_feature_layer_fn = ( - "ferminet.pbc.feature_layer.make_pbc_feature_layer") - cfg.network.make_feature_layer_kwargs = { - "lattice": lattice, - "include_r_ae": False - } - cfg.network.make_envelope_fn = ( - "ferminet.pbc.envelopes.make_multiwave_envelope") - cfg.network.make_envelope_kwargs = {"kpoints": kpoints} - cfg.network.full_det = True return cfg diff --git a/ferminet/jastrows.py b/ferminet/jastrows.py index 23cc83c..b78fda7 100644 --- a/ferminet/jastrows.py +++ b/ferminet/jastrows.py @@ -15,7 +15,7 @@ """Multiplicative Jastrow factors.""" import enum -from typing import Any, Callable, Iterable, Mapping, Union +from typing import Any, Callable, Iterable, Mapping, Union, Optional import jax.numpy as jnp @@ -29,10 +29,42 @@ class JastrowType(enum.Enum): SIMPLE_EE = enum.auto() +def make_periodic_r_ee(lattice: jnp.ndarray): + """Util function for transforming r_ee into the periodic version. + Some of this could be refactored with ferminet.pbc.feature_layer + + Args: + lattice: Matrix whose columns are the primitive lattice vectors of the + system, shape (ndim, ndim). + """ + + # Calculate reciprocal vectors, factor 2pi omitted + reciprocal_vecs = jnp.linalg.inv(lattice) + lattice_metric = lattice.T @ lattice + + def apply(ee: jnp.ndarray): + s_ee = jnp.einsum('il,jkl->jki', reciprocal_vecs, ee) + + n = ee.shape[0] + s_ee += jnp.eye(n)[..., None] + + a = (1 - jnp.cos(2 * jnp.pi * s_ee)) + b = jnp.sin(2 * jnp.pi * s_ee) + cos_term = jnp.einsum('...m,mn,...n->...', a, lattice_metric, a) + sin_term = jnp.einsum('...m,mn,...n->...', b, lattice_metric, b) + periodic_r_ee = (1 / (2 * jnp.pi)) * jnp.sqrt(cos_term + sin_term) + + periodic_r_ee = periodic_r_ee * (1.0 - jnp.eye(n)) + return periodic_r_ee[..., None] + + return apply + + def _jastrow_ee( r_ee: jnp.ndarray, params: ParamTree, nspins: tuple[int, int], + ndim: int, jastrow_fun: Callable[[jnp.ndarray, float, jnp.ndarray], jnp.ndarray], ) -> jnp.ndarray: """Jastrow factor for electron-electron cusps.""" @@ -47,22 +79,29 @@ def _jastrow_ee( if r_ees_parallel.shape[0] > 0: jastrow_ee_par = jnp.sum( - jastrow_fun(r_ees_parallel, 0.25, params['ee_par']) + jastrow_fun(r_ees_parallel, 1 / (ndim + 1), params['ee_par']) ) else: jastrow_ee_par = jnp.asarray(0.0) if r_ees[0][1].shape[0] > 0: - jastrow_ee_anti = jnp.sum(jastrow_fun(r_ees[0][1], 0.5, params['ee_anti'])) + jastrow_ee_anti = jnp.sum( + jastrow_fun(r_ees[0][1], 1 / (ndim - 1), params['ee_anti'])) else: jastrow_ee_anti = jnp.asarray(0.0) return jastrow_ee_anti + jastrow_ee_par -def make_simple_ee_jastrow(): +def make_simple_ee_jastrow(lattice: Optional[jnp.ndarray] = None, ndim: int = 3): """Creates a Jastrow factor for electron-electron cusps.""" + # If working in PBC, use periodic distance for the Jastrow + if lattice is not None: + norm = make_periodic_r_ee(lattice) + else: + norm = lambda x: jnp.linalg.norm(x, axis = -1, keepdims = True) + def simple_ee_cusp_fun( r: jnp.ndarray, cusp: float, alpha: jnp.ndarray ) -> jnp.ndarray: @@ -80,20 +119,24 @@ def init() -> Mapping[str, jnp.ndarray]: return params def apply( - r_ee: jnp.ndarray, + ee: jnp.ndarray, params: ParamTree, nspins: tuple[int, int], ) -> jnp.ndarray: """Jastrow factor for electron-electron cusps.""" - return _jastrow_ee(r_ee, params, nspins, jastrow_fun=simple_ee_cusp_fun) + r_ee = norm(ee) + return _jastrow_ee(r_ee, params, nspins, ndim, jastrow_fun=simple_ee_cusp_fun) return init, apply -def get_jastrow(jastrow: JastrowType): +def get_jastrow( + jastrow: JastrowType, + lattice: Optional[jnp.ndarray] = None, + ndim: int = 3): jastrow_init, jastrow_apply = None, None if jastrow == JastrowType.SIMPLE_EE: - jastrow_init, jastrow_apply = make_simple_ee_jastrow() + jastrow_init, jastrow_apply = make_simple_ee_jastrow(lattice, ndim) elif jastrow != JastrowType.NONE: raise ValueError(f'Unknown Jastrow Factor type: {jastrow}') diff --git a/ferminet/mcmc.py b/ferminet/mcmc.py index e002428..ae92836 100644 --- a/ferminet/mcmc.py +++ b/ferminet/mcmc.py @@ -25,6 +25,20 @@ from jax import lax from jax import numpy as jnp import numpy as np +from typing import Callable + + +def map_to_simulation_cell(pos, lattice, rec, ndim): + """If working in PBC, map electrons back to the supercell + """ + pos_ = jnp.reshape(pos, [-1, ndim]) + phase = jnp.einsum('il,kl->ki', rec / (2 * jnp.pi), pos_) + phase_prim = phase % 1 + prim = jnp.einsum('il,kl->ki', lattice, phase_prim) + return prim.flatten() + +batch_map_to_simulation_cell = jax.vmap( + map_to_simulation_cell, in_axes = (0, None, None, None)) def _harmonic_mean(x, atoms): @@ -78,6 +92,7 @@ def mh_accept(x1, x2, lp_1, lp_2, ratio, key, num_accepts): def mh_update( params: networks.ParamTree, f: networks.LogFermiNetLike, + _map: Callable, data: networks.FermiNetData, key: chex.PRNGKey, lp_1, @@ -94,6 +109,8 @@ def mh_update( params: Wavefuncttion parameters. f: Callable with signature f(params, x) which returns the log of the wavefunction (i.e. the sqaure root of the log probability of x). + _map: If PBC, callable which maps electrons back to the simulation cell. + If OBC, identity function. data: Initial MCMC configurations (batched). key: RNG state. lp_1: log probability of f evaluated at x1 given parameters params. @@ -120,6 +137,7 @@ def mh_update( x1 = data.positions if atoms is None: # symmetric proposal, same stddev everywhere x2 = x1 + stddev * jax.random.normal(subkey, shape=x1.shape) # proposal + x2 = _map(x2) lp_2 = 2.0 * f( params, x2, data.spins, data.atoms, data.charges ) # log prob of proposal @@ -130,6 +148,7 @@ def mh_update( hmean1 = _harmonic_mean(x1, atoms) # harmonic mean of distances to nuclei x2 = x1 + stddev * hmean1 * jax.random.normal(subkey, shape=x1.shape) + x2 = _map(x2) lp_2 = 2.0 * f( params, x2, data.spins, data.atoms, data.charges ) # log prob of proposal @@ -150,6 +169,7 @@ def mh_update( def mh_block_update( params: networks.ParamTree, f: networks.LogFermiNetLike, + _map: Callable, data: networks.FermiNetData, key: chex.PRNGKey, lp_1, @@ -166,6 +186,8 @@ def mh_block_update( params: Wavefuncttion parameters. f: Callable with LogFermiNetLike signature which returns the log of the wavefunction (i.e. the sqaure root of the log probability of x). + _map: If PBC, callable which maps electrons back to the simulation cell. + If OBC, identity function. data: Initial MCMC configuration (batched). key: RNG state. lp_1: log probability of f evaluated at x1 given parameters params. @@ -199,6 +221,7 @@ def mh_block_update( x2 = x1.at[:, ii].add( stddev * jax.random.normal(subkey, shape=x1[:, ii].shape)) x2 = jnp.reshape(x2, [batch_size, -1]) + x2 = _map(x2) if pad > 0: x2 = x2[..., :-pad*ndim] # log prob of proposal @@ -222,7 +245,8 @@ def make_mcmc_step(batch_network, steps=10, atoms=None, ndim=3, - blocks=1): + blocks=1, + lattice=None): """Creates the MCMC step function. Args: @@ -238,12 +262,22 @@ def make_mcmc_step(batch_network, ndim: Dimensionality of the system (usually 3). blocks: Number of blocks to split the updates into. If 1, use all-electron moves. + lattice: If None, assume OBC. Otherwise matrix with supercell lattice + vectors Returns: Callable which performs the set of MCMC steps. """ inner_fun = mh_block_update if blocks > 1 else mh_update + # If PBC, map electrons back to the supercell + if lattice is not None: + rec = 2 * jnp.pi * jnp.linalg.inv(lattice) + _map = lambda p: batch_map_to_simulation_cell(p, lattice, rec, ndim) + else: + rec = None + _map = lambda p: p + def mcmc_step(params, data, key, width): """Performs a set of MCMC steps. @@ -263,6 +297,7 @@ def step_fn(i, x): return inner_fun( params, batch_network, + _map, *x, stddev=width, atoms=atoms, diff --git a/ferminet/networks.py b/ferminet/networks.py index 315d932..9b923aa 100644 --- a/ferminet/networks.py +++ b/ferminet/networks.py @@ -267,6 +267,8 @@ class BaseNetworkOptions: Attributes: ndim: dimension of system. Change only with caution. determinants: Number of determinants to use. + lattice: If None, assume OBC. Otherwise matrix with supercell lattice + vectors. states: Number of outputs, one per excited (or ground) state. Ignored if 0. full_det: If true, evaluate determinants over all electrons. Otherwise, block-diagonalise determinants into spin channels. @@ -281,6 +283,7 @@ class BaseNetworkOptions: """ ndim: int = 3 + lattice: Optional[jnp.ndarray] = None determinants: int = 16 states: int = 0 full_det: bool = True @@ -1079,7 +1082,8 @@ def make_orbitals( equivariant_layers_init, equivariant_layers_apply = equivariant_layers # Optional Jastrow factor. - jastrow_init, jastrow_apply = jastrows.get_jastrow(options.jastrow) + jastrow_init, jastrow_apply = jastrows.get_jastrow( + options.jastrow, options.lattice, options.ndim) def init(key: chex.PRNGKey) -> ParamTree: """Returns initial random parameters for creating orbitals. @@ -1235,7 +1239,7 @@ def apply( # Added pre-determinant for compatibility with pretraining. if jastrow_apply is not None: jastrow = jnp.exp( - jastrow_apply(r_ee, params['jastrow'], nspins) / sum(nspins) + jastrow_apply(ee, params['jastrow'], nspins) / sum(nspins) ) orbitals = [orbital * jastrow for orbital in orbitals] @@ -1365,6 +1369,7 @@ def make_fermi_net( charges: jnp.ndarray, *, ndim: int = 3, + lattice: Optional[jnp.ndarray] = None, determinants: int = 16, states: int = 0, envelope: Optional[envelopes.Envelope] = None, @@ -1389,6 +1394,8 @@ def make_fermi_net( nspins: Tuple of the number of spin-up and spin-down electrons. charges: (natom) array of atom nuclear charges. ndim: dimension of system. Change only with caution. + lattice: If None, assume OBC. Otherwise matrix with supercell lattice + vectors. determinants: Number of determinants to use. states: Number of outputs, one per excited (or ground) state. Ignored if 0. envelope: Envelope to use to impose orbitals go to zero at infinity. @@ -1448,6 +1455,7 @@ def make_fermi_net( options = FermiNetOptions( ndim=ndim, + lattice=lattice, determinants=determinants, states=states, rescale_inputs=rescale_inputs, diff --git a/ferminet/pbc/envelopes.py b/ferminet/pbc/envelopes.py index 540b712..bbaf6a6 100644 --- a/ferminet/pbc/envelopes.py +++ b/ferminet/pbc/envelopes.py @@ -43,8 +43,7 @@ def make_multiwave_envelope(kpoints: jnp.ndarray) -> envelopes.Envelope: Args: kpoints: Reciprocal lattice vectors of terms included in the Fourier - series. Shape (nkpoints, ndim) (Note that ndim=3 is currently - a hard-coded default). + series. Shape (nkpoints, ndim) Returns: An instance of ferminet.envelopes.Envelope with apply_type @@ -79,6 +78,7 @@ def apply(*, ae: jnp.ndarray, r_ae: jnp.ndarray, r_ee: jnp.ndarray, def make_kpoints( lattice: Union[np.ndarray, jnp.ndarray], spins: Tuple[int, int], + ndim: int = 3, min_kpoints: Optional[int] = None, ) -> jnp.ndarray: """Generates an array of reciprocal lattice vectors. @@ -88,6 +88,7 @@ def make_kpoints( system, shape (ndim, ndim). (Note that ndim=3 is currently a hard-coded default). spins: Tuple of the number of spin-up and spin-down electrons. + ndim: Number of dimensions. min_kpoints: If specified, the number of kpoints which must be included in the output. The number of kpoints returned will be the first filled shell which is larger than this value. Defaults to None, @@ -111,9 +112,9 @@ def make_kpoints( dk = 1 + 1e-5 # Generate ordinals of the lowest min_kpoints kpoints - max_k = int(jnp.ceil(min_kpoints * dk)**(1 / 3.)) + max_k = int(jnp.ceil(min_kpoints * dk)**(1 / ndim)) ordinals = sorted(range(-max_k, max_k+1), key=abs) - ordinals = jnp.asarray(list(itertools.product(ordinals, repeat=3))) + ordinals = jnp.asarray(list(itertools.product(ordinals, repeat=ndim))) kpoints = ordinals @ rec_lattice.T kpoints = jnp.asarray(sorted(kpoints, key=jnp.linalg.norm)) diff --git a/ferminet/pbc/ewald2d.py b/ferminet/pbc/ewald2d.py new file mode 100644 index 0000000..7b79126 --- /dev/null +++ b/ferminet/pbc/ewald2d.py @@ -0,0 +1,146 @@ +# Copyright 2022 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License + +"""2D Ewald summation of Coulomb Hamiltonian in periodic boundary conditions. +""" + +import itertools +from typing import Callable + +import jax +import jax.numpy as jnp + + +def make_ewald_potential( + lattice: jnp.ndarray, + atoms: jnp.ndarray, + charges: jnp.ndarray, + truncation_limit: int = 5, + include_heg_background: bool = True +) -> Callable[[jnp.ndarray, jnp.ndarray], float]: + """Creates a function to evaluate infinite Coulomb sum for periodic lattice. + + Args: + lattice: Shape (2, 2). Matrix whose columns are the primitive lattice + vectors. + atoms: Shape (natoms, ndim). Positions of the atoms. + charges: Shape (natoms). Nuclear charges of the atoms. + truncation_limit: Integer. Half side length of cube of nearest neighbours + to primitive cell which are summed over in evaluation of Ewald sum. + Must be large enough to achieve convergence for the real and reciprocal + space sums. + include_heg_background: bool. When True, includes cell-neutralizing + background term for homogeneous electron gas. + + Returns: + Callable with signature f(ae, ee), where (ae, ee) are atom-electon and + electron-electron displacement vectors respectively, which evaluates the + Coulomb sum for the periodic lattice via the Ewald method. + """ + rec = 2 * jnp.pi * jnp.linalg.inv(lattice) + volume = jnp.abs(jnp.linalg.det(lattice)) + # the factor gamma tunes the width of the summands in real / reciprocal space + # and this value is chosen to optimize the convergence trade-off between the + # two sums. See CASINO QMC manual. + root_gamma = 2.4 / volume**0.5 + ordinals = sorted(range(-truncation_limit, truncation_limit + 1), key=abs) + ordinals = jnp.array(list(itertools.product(ordinals, repeat=2))) + lat_vectors = jnp.einsum('kj,ij->ik', lattice, ordinals) + rec_vectors = jnp.einsum('jk,ij->ik', rec, ordinals[1:]) + rec_vec_square = jnp.einsum('ij,ij->i', rec_vectors, rec_vectors) + rec_vec_norm = jnp.sqrt(rec_vec_square) + lat_vec_norm = jnp.linalg.norm(lat_vectors[1:], axis=-1) + + def real_space_ewald(separation: jnp.ndarray): + """Real-space Ewald potential between charges seperated by separation.""" + displacements = jnp.linalg.norm( + separation - lat_vectors, axis=-1) # |r - R| + return jnp.sum( + jax.scipy.special.erfc(root_gamma * displacements) / displacements) + + def recp_space_ewald(separation: jnp.ndarray): + """Returns reciprocal-space Ewald potential between charges.""" + return (2 * jnp.pi / volume) * jnp.sum( + jnp.exp(1.0j * jnp.dot(rec_vectors, separation)) * + jax.scipy.special.erfc(rec_vec_norm / (2 * root_gamma)) / rec_vec_norm) + + def ewald_sum(separation: jnp.ndarray): + """Evaluates combined real and reciprocal space Ewald potential.""" + return (real_space_ewald(separation) + recp_space_ewald(separation) - + 2 * (jnp.pi**0.5) / (volume * root_gamma)) + + madelung_const = ( + jnp.sum(jax.scipy.special.erfc(root_gamma * lat_vec_norm) / lat_vec_norm) + - 2 * root_gamma / jnp.pi**0.5) + madelung_const += ( + (2 * jnp.pi / volume) * + jnp.sum(jax.scipy.special.erfc(rec_vec_norm / (2 * root_gamma)) / rec_vec_norm) - + 2 * (jnp.pi**0.5) / (volume * root_gamma)) + + batch_ewald_sum = jax.vmap(ewald_sum, in_axes=(0,)) + + def atom_electron_potential(ae: jnp.ndarray): + """Evaluates periodic atom-electron potential.""" + nelec = ae.shape[0] + ae = jnp.reshape(ae, [-1, 2]) # flatten electronxatom axis + # calculate potential for each ae pair + ewald = batch_ewald_sum(ae) - madelung_const + return jnp.sum(-jnp.tile(charges, nelec) * ewald) + + def electron_electron_potential(ee: jnp.ndarray): + """Evaluates periodic electron-electron potential.""" + nelec = ee.shape[0] + ee = jnp.reshape(ee, [-1, 2]) + if include_heg_background: + ewald = batch_ewald_sum(ee) + else: + ewald = batch_ewald_sum(ee) - madelung_const + ewald = jnp.reshape(ewald, [nelec, nelec]) + ewald = ewald.at[jnp.diag_indices(nelec)].set(0.0) + if include_heg_background: + return 0.5 * jnp.sum(ewald) + 0.5 * nelec * madelung_const + else: + return 0.5 * jnp.sum(ewald) + + # Atom-atom potential + natom = atoms.shape[0] + if natom > 1: + aa = jnp.reshape(atoms, [1, -1, 2]) - jnp.reshape(atoms, [-1, 1, 2]) + aa = jnp.reshape(aa, [-1, 2]) + chargeprods = (charges[..., None] @ charges[..., None].T).flatten() + ewald = batch_ewald_sum(aa) - madelung_const + ewald = jnp.reshape(ewald, [natom, natom]) + ewald = ewald.at[jnp.diag_indices(natom)].set(0.0) + ewald = ewald.flatten() + atom_atom_potential = 0.5 * jnp.sum(chargeprods * ewald) + else: + atom_atom_potential = 0.0 + + def potential(ae: jnp.ndarray, ee: jnp.ndarray): + """Accumulates atom-electron, atom-atom, and electron-electron potential.""" + # Reduce vectors into first unit cell - Ewald summation + # is only guaranteed to converge close to the origin + """ Should not be needed if mcmc_pbc is used + phase_ae = jnp.einsum('il,jkl->jki', rec / (2 * jnp.pi), ae) + phase_ee = jnp.einsum('il,jkl->jki', rec / (2 * jnp.pi), ee) + phase_prim_ae = phase_ae % 1 + phase_prim_ee = phase_ee % 1 + prim_ae = jnp.einsum('il,jkl->jki', lattice, phase_prim_ae) + prim_ee = jnp.einsum('il,jkl->jki', lattice, phase_prim_ee) + """ + return jnp.real( + atom_electron_potential(ae) + # prim_ae + electron_electron_potential(ee) + atom_atom_potential) # prim_ee + + return potential diff --git a/ferminet/pbc/ewald3d.py b/ferminet/pbc/ewald3d.py new file mode 100644 index 0000000..c00e785 --- /dev/null +++ b/ferminet/pbc/ewald3d.py @@ -0,0 +1,147 @@ +# Copyright 2022 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License + +"""3D Ewald summation of Coulomb Hamiltonian in periodic boundary conditions. + +See Cassella, G., Sutterud, H., Azadi, S., Drummond, N.D., Pfau, D., +Spencer, J.S. and Foulkes, W.M.C., 2022. Discovering Quantum Phase Transitions +with Fermionic Neural Networks. arXiv preprint arXiv:2202.05183. +""" + +import itertools +from typing import Callable + +import jax +import jax.numpy as jnp + + +def make_ewald_potential( + lattice: jnp.ndarray, + atoms: jnp.ndarray, + charges: jnp.ndarray, + truncation_limit: int = 5, + include_heg_background: bool = True +) -> Callable[[jnp.ndarray, jnp.ndarray], float]: + """Creates a function to evaluate infinite Coulomb sum for periodic lattice. + + Args: + lattice: Shape (3, 3). Matrix whose columns are the primitive lattice + vectors. + atoms: Shape (natoms, ndim). Positions of the atoms. + charges: Shape (natoms). Nuclear charges of the atoms. + truncation_limit: Integer. Half side length of cube of nearest neighbours + to primitive cell which are summed over in evaluation of Ewald sum. + Must be large enough to achieve convergence for the real and reciprocal + space sums. + include_heg_background: bool. When True, includes cell-neutralizing + background term for homogeneous electron gas. + + Returns: + Callable with signature f(ae, ee), where (ae, ee) are atom-electon and + electron-electron displacement vectors respectively, which evaluates the + Coulomb sum for the periodic lattice via the Ewald method. + """ + rec = 2 * jnp.pi * jnp.linalg.inv(lattice) + volume = jnp.abs(jnp.linalg.det(lattice)) + # the factor gamma tunes the width of the summands in real / reciprocal space + # and this value is chosen to optimize the convergence trade-off between the + # two sums. See CASINO QMC manual. + gamma = (2.8 / volume**(1 / 3))**2 + ordinals = sorted(range(-truncation_limit, truncation_limit + 1), key=abs) + ordinals = jnp.array(list(itertools.product(ordinals, repeat=3))) + lat_vectors = jnp.einsum('kj,ij->ik', lattice, ordinals) + rec_vectors = jnp.einsum('jk,ij->ik', rec, ordinals[1:]) + rec_vec_square = jnp.einsum('ij,ij->i', rec_vectors, rec_vectors) + lat_vec_norm = jnp.linalg.norm(lat_vectors[1:], axis=-1) + + def real_space_ewald(separation: jnp.ndarray): + """Real-space Ewald potential between charges seperated by separation.""" + displacements = jnp.linalg.norm( + separation - lat_vectors, axis=-1) # |r - R| + return jnp.sum( + jax.scipy.special.erfc(gamma**0.5 * displacements) / displacements) + + def recp_space_ewald(separation: jnp.ndarray): + """Returns reciprocal-space Ewald potential between charges.""" + return (4 * jnp.pi / volume) * jnp.sum( + jnp.exp(1.0j * jnp.dot(rec_vectors, separation)) * + jnp.exp(-rec_vec_square / (4 * gamma)) / rec_vec_square) + + def ewald_sum(separation: jnp.ndarray): + """Evaluates combined real and reciprocal space Ewald potential.""" + return (real_space_ewald(separation) + recp_space_ewald(separation) - + jnp.pi / (volume * gamma)) + + madelung_const = ( + jnp.sum(jax.scipy.special.erfc(gamma**0.5 * lat_vec_norm) / lat_vec_norm) + - 2 * gamma**0.5 / jnp.pi**0.5) + madelung_const += ( + (4 * jnp.pi / volume) * + jnp.sum(jnp.exp(-rec_vec_square / (4 * gamma)) / rec_vec_square) - + jnp.pi / (volume * gamma)) + + batch_ewald_sum = jax.vmap(ewald_sum, in_axes=(0,)) + + def atom_electron_potential(ae: jnp.ndarray): + """Evaluates periodic atom-electron potential.""" + nelec = ae.shape[0] + ae = jnp.reshape(ae, [-1, 3]) # flatten electronxatom axis + # calculate potential for each ae pair + ewald = batch_ewald_sum(ae) - madelung_const + return jnp.sum(-jnp.tile(charges, nelec) * ewald) + + def electron_electron_potential(ee: jnp.ndarray): + """Evaluates periodic electron-electron potential.""" + nelec = ee.shape[0] + ee = jnp.reshape(ee, [-1, 3]) + if include_heg_background: + ewald = batch_ewald_sum(ee) + else: + ewald = batch_ewald_sum(ee) - madelung_const + ewald = jnp.reshape(ewald, [nelec, nelec]) + ewald = ewald.at[jnp.diag_indices(nelec)].set(0.0) + if include_heg_background: + return 0.5 * jnp.sum(ewald) + 0.5 * nelec * madelung_const + else: + return 0.5 * jnp.sum(ewald) + + # Atom-atom potential + natom = atoms.shape[0] + if natom > 1: + aa = jnp.reshape(atoms, [1, -1, 3]) - jnp.reshape(atoms, [-1, 1, 3]) + aa = jnp.reshape(aa, [-1, 3]) + chargeprods = (charges[..., None] @ charges[..., None].T).flatten() + ewald = batch_ewald_sum(aa) - madelung_const + ewald = jnp.reshape(ewald, [natom, natom]) + ewald = ewald.at[jnp.diag_indices(natom)].set(0.0) + ewald = ewald.flatten() + atom_atom_potential = 0.5 * jnp.sum(chargeprods * ewald) + else: + atom_atom_potential = 0.0 + + def potential(ae: jnp.ndarray, ee: jnp.ndarray): + """Accumulates atom-electron, atom-atom, and electron-electron potential.""" + # Reduce vectors into first unit cell - Ewald summation + # is only guaranteed to converge close to the origin + phase_ae = jnp.einsum('il,jkl->jki', rec / (2 * jnp.pi), ae) + phase_ee = jnp.einsum('il,jkl->jki', rec / (2 * jnp.pi), ee) + phase_prim_ae = phase_ae % 1 + phase_prim_ee = phase_ee % 1 + prim_ae = jnp.einsum('il,jkl->jki', lattice, phase_prim_ae) + prim_ee = jnp.einsum('il,jkl->jki', lattice, phase_prim_ee) + return jnp.real( + atom_electron_potential(prim_ae) + + electron_electron_potential(prim_ee) + atom_atom_potential) + + return potential diff --git a/ferminet/pbc/feature_layer.py b/ferminet/pbc/feature_layer.py index f926a71..6b8679d 100644 --- a/ferminet/pbc/feature_layer.py +++ b/ferminet/pbc/feature_layer.py @@ -62,7 +62,8 @@ def make_pbc_feature_layer( lattice: Matrix whose columns are the primitive lattice vectors of the system, shape (ndim, ndim). include_r_ae: Flag to enable electron-atom distance features. Set to False - to avoid cusps with ghost atoms in, e.g., homogeneous electron gas. + to avoid cusps with ghost atoms in, e.g., homogeneous electron gas. If not + needed, set cfg.network.make_feature_layer_kwargs['include_r_ae'] = False """ del nspins diff --git a/ferminet/pbc/hamiltonian.py b/ferminet/pbc/hamiltonian.py index 07d1eba..7c142a3 100644 --- a/ferminet/pbc/hamiltonian.py +++ b/ferminet/pbc/hamiltonian.py @@ -19,142 +19,22 @@ with Fermionic Neural Networks. arXiv preprint arXiv:2202.05183. """ -import itertools from typing import Callable, Optional, Sequence, Tuple import chex from ferminet import hamiltonian from ferminet import networks -import jax +from ferminet.pbc.ewald2d import make_ewald_potential as ewald2d +from ferminet.pbc.ewald3d import make_ewald_potential as ewald3d import jax.numpy as jnp -def make_ewald_potential( - lattice: jnp.ndarray, - atoms: jnp.ndarray, - charges: jnp.ndarray, - truncation_limit: int = 5, - include_heg_background: bool = True -) -> Callable[[jnp.ndarray, jnp.ndarray], float]: - """Creates a function to evaluate infinite Coulomb sum for periodic lattice. - - Args: - lattice: Shape (3, 3). Matrix whose columns are the primitive lattice - vectors. - atoms: Shape (natoms, ndim). Positions of the atoms. - charges: Shape (natoms). Nuclear charges of the atoms. - truncation_limit: Integer. Half side length of cube of nearest neighbours - to primitive cell which are summed over in evaluation of Ewald sum. - Must be large enough to achieve convergence for the real and reciprocal - space sums. - include_heg_background: bool. When True, includes cell-neutralizing - background term for homogeneous electron gas. - - Returns: - Callable with signature f(ae, ee), where (ae, ee) are atom-electon and - electron-electron displacement vectors respectively, which evaluates the - Coulomb sum for the periodic lattice via the Ewald method. - """ - rec = 2 * jnp.pi * jnp.linalg.inv(lattice) - volume = jnp.abs(jnp.linalg.det(lattice)) - # the factor gamma tunes the width of the summands in real / reciprocal space - # and this value is chosen to optimize the convergence trade-off between the - # two sums. See CASINO QMC manual. - gamma = (2.8 / volume**(1 / 3))**2 - ordinals = sorted(range(-truncation_limit, truncation_limit + 1), key=abs) - ordinals = jnp.array(list(itertools.product(ordinals, repeat=3))) - lat_vectors = jnp.einsum('kj,ij->ik', lattice, ordinals) - rec_vectors = jnp.einsum('jk,ij->ik', rec, ordinals[1:]) - rec_vec_square = jnp.einsum('ij,ij->i', rec_vectors, rec_vectors) - lat_vec_norm = jnp.linalg.norm(lat_vectors[1:], axis=-1) - - def real_space_ewald(separation: jnp.ndarray): - """Real-space Ewald potential between charges seperated by separation.""" - displacements = jnp.linalg.norm( - separation - lat_vectors, axis=-1) # |r - R| - return jnp.sum( - jax.scipy.special.erfc(gamma**0.5 * displacements) / displacements) - - def recp_space_ewald(separation: jnp.ndarray): - """Returns reciprocal-space Ewald potential between charges.""" - return (4 * jnp.pi / volume) * jnp.sum( - jnp.exp(1.0j * jnp.dot(rec_vectors, separation)) * - jnp.exp(-rec_vec_square / (4 * gamma)) / rec_vec_square) - - def ewald_sum(separation: jnp.ndarray): - """Evaluates combined real and reciprocal space Ewald potential.""" - return (real_space_ewald(separation) + recp_space_ewald(separation) - - jnp.pi / (volume * gamma)) - - madelung_const = ( - jnp.sum(jax.scipy.special.erfc(gamma**0.5 * lat_vec_norm) / lat_vec_norm) - - 2 * gamma**0.5 / jnp.pi**0.5) - madelung_const += ( - (4 * jnp.pi / volume) * - jnp.sum(jnp.exp(-rec_vec_square / (4 * gamma)) / rec_vec_square) - - jnp.pi / (volume * gamma)) - - batch_ewald_sum = jax.vmap(ewald_sum, in_axes=(0,)) - - def atom_electron_potential(ae: jnp.ndarray): - """Evaluates periodic atom-electron potential.""" - nelec = ae.shape[0] - ae = jnp.reshape(ae, [-1, 3]) # flatten electronxatom axis - # calculate potential for each ae pair - ewald = batch_ewald_sum(ae) - madelung_const - return jnp.sum(-jnp.tile(charges, nelec) * ewald) - - def electron_electron_potential(ee: jnp.ndarray): - """Evaluates periodic electron-electron potential.""" - nelec = ee.shape[0] - ee = jnp.reshape(ee, [-1, 3]) - if include_heg_background: - ewald = batch_ewald_sum(ee) - else: - ewald = batch_ewald_sum(ee) - madelung_const - ewald = jnp.reshape(ewald, [nelec, nelec]) - ewald = ewald.at[jnp.diag_indices(nelec)].set(0.0) - if include_heg_background: - return 0.5 * jnp.sum(ewald) + 0.5 * nelec * madelung_const - else: - return 0.5 * jnp.sum(ewald) - - # Atom-atom potential - natom = atoms.shape[0] - if natom > 1: - aa = jnp.reshape(atoms, [1, -1, 3]) - jnp.reshape(atoms, [-1, 1, 3]) - aa = jnp.reshape(aa, [-1, 3]) - chargeprods = (charges[..., None] @ charges[..., None].T).flatten() - ewald = batch_ewald_sum(aa) - madelung_const - ewald = jnp.reshape(ewald, [natom, natom]) - ewald = ewald.at[jnp.diag_indices(natom)].set(0.0) - ewald = ewald.flatten() - atom_atom_potential = 0.5 * jnp.sum(chargeprods * ewald) - else: - atom_atom_potential = 0.0 - - def potential(ae: jnp.ndarray, ee: jnp.ndarray): - """Accumulates atom-electron, atom-atom, and electron-electron potential.""" - # Reduce vectors into first unit cell - Ewald summation - # is only guaranteed to converge close to the origin - phase_ae = jnp.einsum('il,jkl->jki', rec / (2 * jnp.pi), ae) - phase_ee = jnp.einsum('il,jkl->jki', rec / (2 * jnp.pi), ee) - phase_prim_ae = phase_ae % 1 - phase_prim_ee = phase_ee % 1 - prim_ae = jnp.einsum('il,jkl->jki', lattice, phase_prim_ae) - prim_ee = jnp.einsum('il,jkl->jki', lattice, phase_prim_ee) - return jnp.real( - atom_electron_potential(prim_ae) + - electron_electron_potential(prim_ee) + atom_atom_potential) - - return potential - - def local_energy( f: networks.FermiNetLike, charges: jnp.ndarray, nspins: Sequence[int], use_scan: bool = False, + ndim: int = 3, complex_output: bool = False, laplacian_method: str = 'default', states: int = 0, @@ -162,7 +42,7 @@ def local_energy( pp_type: str = 'ccecp', pp_symbols: Sequence[str] | None = None, lattice: Optional[jnp.ndarray] = None, - heg: bool = True, + heg: bool = False, convergence_radius: int = 5, ) -> hamiltonian.LocalEnergy: """Creates the local energy function in periodic boundary conditions. @@ -173,6 +53,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 @@ -186,6 +67,8 @@ def local_energy( lattice: Shape (ndim, ndim). Matrix of lattice vectors. Default: identity matrix. heg: bool. Flag to enable features specific to the electron gas. + This is False by default, if needed, set + cfg.system.make_local_energy_kwargs['heg'] = True convergence_radius: int. Radius of cluster summed over by Ewald sums. Returns: @@ -200,13 +83,18 @@ def local_energy( del nspins del pp_type - if lattice is None: - lattice = jnp.eye(3) ke = hamiltonian.local_kinetic_energy(f, use_scan=use_scan, complex_output=complex_output, laplacian_method=laplacian_method) + + if ndim == 3: + make_potential_energy = ewald3d + elif ndim == 2: + make_potential_energy = ewald2d + else: + raise NotImplementedError(f"Ewald sum for ndim = {ndim} not implemented") def _e_l( params: networks.ParamTree, key: chex.PRNGKey, data: networks.FermiNetData @@ -219,11 +107,11 @@ def _e_l( data: MCMC configuration. """ del key # unused - potential_energy = make_ewald_potential( + potential_energy = make_potential_energy( lattice, data.atoms, charges, convergence_radius, heg ) ae, ee, _, _ = networks.construct_input_features( - data.positions, data.atoms) + data.positions, data.atoms, ndim) potential = potential_energy(ae, ee) kinetic = ke(params, data) return potential + kinetic, None diff --git a/ferminet/psiformer.py b/ferminet/psiformer.py index a897616..4470621 100644 --- a/ferminet/psiformer.py +++ b/ferminet/psiformer.py @@ -337,6 +337,7 @@ def make_fermi_net( charges: jnp.ndarray, *, ndim: int = 3, + lattice: Optional[jnp.ndarray] = None, determinants: int = 16, states: int = 0, envelope: Optional[envelopes.Envelope] = None, @@ -361,6 +362,8 @@ def make_fermi_net( nspins: Tuple of the number of spin-up and spin-down electrons. charges: (natom) array of atom nuclear charges. ndim: Dimension of the system. Change only with caution. + lattice: If None, assume OBC. Otherwise matrix with supercell lattice + vectors. determinants: Number of determinants. states: Number of outputs, one per excited (or ground) state. Ignored if 0. envelope: Envelope to use to impose orbitals go to zero at infinity. @@ -401,6 +404,7 @@ def make_fermi_net( options = PsiformerOptions( ndim=ndim, + lattice=lattice, determinants=determinants, states=states, envelope=envelope, diff --git a/ferminet/train.py b/ferminet/train.py index 41333af..2cdc3d2 100644 --- a/ferminet/train.py +++ b/ferminet/train.py @@ -33,6 +33,9 @@ from ferminet import observables from ferminet import pretrain from ferminet import psiformer +from ferminet.pbc import hamiltonian as pbc_hamiltonian +from ferminet.pbc import feature_layer as pbc_feature_layer +from ferminet.pbc import envelopes as pbc_envelopes from ferminet.utils import statistics from ferminet.utils import system from ferminet.utils import utils @@ -481,18 +484,18 @@ def train(cfg: ml_collections.ConfigDict, writer_manager=None): make_feature_layer: networks.MakeFeatureLayer = getattr( feature_layer_module, feature_layer_fn ) - feature_layer = make_feature_layer( - natoms=charges.shape[0], - nspins=cfg.system.electrons, - ndim=cfg.system.ndim, - **cfg.network.make_feature_layer_kwargs) + elif cfg.system.lattice is not None: + make_feature_layer = pbc_feature_layer.make_pbc_feature_layer + cfg.network.make_feature_layer_kwargs['lattice'] = cfg.system.lattice else: - feature_layer = networks.make_ferminet_features( - natoms=charges.shape[0], - nspins=cfg.system.electrons, - ndim=cfg.system.ndim, - rescale_inputs=cfg.network.get('rescale_inputs', False), - ) + make_feature_layer = networks.make_ferminet_features + feature_layer = make_feature_layer( + natoms=charges.shape[0], + nspins=cfg.system.electrons, + ndim=cfg.system.ndim, + rescale_inputs=cfg.network.get('rescale_inputs', False), + **cfg.network.make_feature_layer_kwargs + ) if cfg.network.make_envelope_fn: envelope_module, envelope_fn = ( @@ -500,6 +503,10 @@ def train(cfg: ml_collections.ConfigDict, writer_manager=None): envelope_module = importlib.import_module(envelope_module) make_envelope = getattr(envelope_module, envelope_fn) envelope = make_envelope(**cfg.network.make_envelope_kwargs) # type: envelopes.Envelope + elif cfg.system.lattice is not None: + kpoints = pbc_envelopes.make_kpoints( + cfg.system.lattice, cfg.system.electrons, cfg.system.ndim) + envelope = pbc_envelopes.make_multiwave_envelope(kpoints) else: envelope = envelopes.make_isotropic_envelope() @@ -509,6 +516,7 @@ def train(cfg: ml_collections.ConfigDict, writer_manager=None): nspins, charges, ndim=cfg.system.ndim, + lattice=cfg.system.lattice, determinants=cfg.network.determinants, states=cfg.system.states, envelope=envelope, @@ -525,6 +533,7 @@ def train(cfg: ml_collections.ConfigDict, writer_manager=None): nspins, charges, ndim=cfg.system.ndim, + lattice=cfg.system.lattice, determinants=cfg.network.determinants, states=cfg.system.states, envelope=envelope, @@ -740,6 +749,7 @@ def log_network(*args, **kwargs): atoms=atoms_to_mcmc, blocks=cfg.mcmc.blocks * num_states, ndim=cfg.system.ndim, + lattice=cfg.system.lattice, ) # Construct loss and optimizer @@ -748,6 +758,9 @@ def log_network(*args, **kwargs): 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 + elif cfg.system.lattice is not None: + make_local_energy = pbc_hamiltonian.local_energy + cfg.system.make_local_energy_kwargs['lattice'] = cfg.system.lattice else: make_local_energy = hamiltonian.local_energy laplacian_method = cfg.optim.get('laplacian', 'default')