216 lines
8.0 KiB
Python
216 lines
8.0 KiB
Python
"""
|
||
Simulation cell with plane-wave basis set.
|
||
|
||
Convention (Hartree atomic units throughout):
|
||
f(r) = (1/V) × Σ_G f̂(G) exp(i G·r)
|
||
f̂(G) = ∫_V f(r) exp(-i G·r) dr ≈ dV × FFT[f]_n
|
||
|
||
Discrete form (numpy FFT):
|
||
FFT[f]_n = Σ_j f(r_j) exp(-2πi j·n/N) ← numpy.fft.fftn
|
||
f(r_j) = ifftn(F)[j] where F_n = f̂(G_n)/dV ← numpy.fft.ifftn
|
||
|
||
So: G→r: f(r) = ifftn(f̂/dV).ravel()
|
||
r→G: f̂(G) = dV × fftn(f)[n]
|
||
"""
|
||
|
||
import numpy as np
|
||
|
||
|
||
class Cell:
|
||
"""
|
||
Periodic simulation cell with a plane-wave basis.
|
||
|
||
Parameters
|
||
----------
|
||
a_matrix : (3,3) array
|
||
Rows are lattice vectors in Bohr.
|
||
ecut : float
|
||
Wavefunction kinetic-energy cutoff in Hartree.
|
||
kpt : (3,) array, optional
|
||
k-point in Cartesian reciprocal space (default: Γ).
|
||
"""
|
||
|
||
def __init__(self, a_matrix, ecut: float, kpt=None):
|
||
self.a = np.asarray(a_matrix, dtype=float)
|
||
self.ecut = float(ecut)
|
||
self.kpt = np.zeros(3) if kpt is None else np.asarray(kpt, dtype=float)
|
||
|
||
self.volume = abs(np.linalg.det(self.a))
|
||
self.b = 2 * np.pi * np.linalg.inv(self.a).T # rows are b_i
|
||
|
||
self._build_gvecs()
|
||
self._build_grid()
|
||
self._build_full_G2()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Setup
|
||
# ------------------------------------------------------------------
|
||
|
||
def _build_gvecs(self):
|
||
"""G-vectors with |k+G|²/2 ≤ ecut (wavefunction cutoff)."""
|
||
b_lens = np.linalg.norm(self.b, axis=1)
|
||
n_max = np.ceil(np.sqrt(2 * self.ecut) / b_lens).astype(int) + 1
|
||
|
||
ns = [np.arange(-n, n + 1) for n in n_max]
|
||
n1, n2, n3 = np.meshgrid(*ns, indexing='ij')
|
||
n_all = np.stack([n1.ravel(), n2.ravel(), n3.ravel()], axis=1)
|
||
|
||
G = n_all @ self.b
|
||
kG = G + self.kpt
|
||
KE = 0.5 * np.einsum('ij,ij->i', kG, kG)
|
||
|
||
mask = KE <= self.ecut
|
||
self.G = G[mask]
|
||
self.G_int = n_all[mask]
|
||
self.KE_G = KE[mask]
|
||
self.N_G = int(mask.sum())
|
||
|
||
def _build_grid(self):
|
||
"""Real-space FFT grid, large enough for the density cutoff (4×ecut)."""
|
||
b_lens = np.linalg.norm(self.b, axis=1)
|
||
n_dens = np.ceil(2 * np.sqrt(2 * self.ecut) / b_lens).astype(int)
|
||
self.mesh = 2 * n_dens + 1 # odd → symmetric
|
||
self.N_r = int(np.prod(self.mesh))
|
||
self.dV = self.volume / self.N_r
|
||
|
||
s = [np.arange(m) / m for m in self.mesh]
|
||
i0, i1, i2 = np.meshgrid(*s, indexing='ij')
|
||
frac = np.stack([i0.ravel(), i1.ravel(), i2.ravel()], axis=1)
|
||
self.r = frac @ self.a # (N_r, 3) in Bohr
|
||
|
||
def _build_full_G2(self):
|
||
"""Precompute G-vectors and |G|² for every point on the full FFT mesh."""
|
||
def freqs(n):
|
||
return np.fft.fftfreq(n) * n # integer frequency indices
|
||
|
||
n0 = freqs(self.mesh[0])
|
||
n1 = freqs(self.mesh[1])
|
||
n2 = freqs(self.mesh[2])
|
||
i0, i1, i2 = np.meshgrid(n0, n1, n2, indexing='ij')
|
||
# G_full[i,j,k] = i*b0 + j*b1 + k*b2
|
||
Gx = i0 * self.b[0, 0] + i1 * self.b[1, 0] + i2 * self.b[2, 0]
|
||
Gy = i0 * self.b[0, 1] + i1 * self.b[1, 1] + i2 * self.b[2, 1]
|
||
Gz = i0 * self.b[0, 2] + i1 * self.b[1, 2] + i2 * self.b[2, 2]
|
||
self._G2_full = Gx**2 + Gy**2 + Gz**2 # (N0,N1,N2)
|
||
# Flat (N_r, 3) G-vectors for structure factor computation
|
||
self.G_full = np.stack([Gx.ravel(), Gy.ravel(), Gz.ravel()], axis=1)
|
||
|
||
# ------------------------------------------------------------------
|
||
# FFT core — keep shapes flat (N_r,) for real space
|
||
# ------------------------------------------------------------------
|
||
|
||
def _wf_idx(self):
|
||
"""Index arrays to address wavefunction G-vectors on the full mesh."""
|
||
return tuple((self.G_int[:, d] % self.mesh[d]) for d in range(3))
|
||
|
||
def G_to_r(self, f_G: np.ndarray) -> np.ndarray:
|
||
"""
|
||
G-space → real-space. Returns flat (N_r,) float array.
|
||
f_G holds physical f̂(G) coefficients.
|
||
"""
|
||
F_full = np.zeros(self.mesh, dtype=complex)
|
||
ix, iy, iz = self._wf_idx()
|
||
F_full[ix, iy, iz] = f_G / self.dV # F_n = f̂_n / dV
|
||
return np.fft.ifftn(F_full).real.ravel()
|
||
|
||
def r_to_G(self, f_r: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Real-space → G-space. Returns complex (N_G,) array (wavefunction cutoff).
|
||
f_r is a flat (N_r,) real array.
|
||
"""
|
||
F_full = np.fft.fftn(f_r.reshape(self.mesh))
|
||
ix, iy, iz = self._wf_idx()
|
||
return F_full[ix, iy, iz] * self.dV # f̂_G = dV × FFT_n
|
||
|
||
# ------------------------------------------------------------------
|
||
# Density
|
||
# ------------------------------------------------------------------
|
||
|
||
def density_from_orbitals(self, orbs_G: np.ndarray, occ: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Electron density on the real-space grid.
|
||
|
||
Parameters
|
||
----------
|
||
orbs_G : (n_occ, N_G) complex array — orbital G-coefficients.
|
||
occ : (n_occ,) float array — occupation numbers.
|
||
|
||
Returns
|
||
-------
|
||
rho_r : (N_r,) float array e/Bohr³.
|
||
"""
|
||
rho_r = np.zeros(self.N_r)
|
||
for n, f in enumerate(occ):
|
||
if f == 0.0:
|
||
continue
|
||
psi_r = self.G_to_r(orbs_G[n]) # (N_r,)
|
||
rho_r += f * psi_r ** 2
|
||
return rho_r
|
||
|
||
# ------------------------------------------------------------------
|
||
# Hartree potential (Poisson on the full density-cutoff mesh)
|
||
# ------------------------------------------------------------------
|
||
|
||
def hartree_potential(self, rho_r: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Solve ∇²V_H = -4π ρ via FFT on the full mesh.
|
||
Returns V_H as a flat (N_r,) float array.
|
||
"""
|
||
rho_G_full = np.fft.fftn(rho_r.reshape(self.mesh)) * self.dV # f̂(G)
|
||
|
||
G2 = self._G2_full
|
||
V_H_G_full = np.zeros_like(rho_G_full)
|
||
nz = G2 > 1e-12
|
||
V_H_G_full[nz] = 4 * np.pi * rho_G_full[nz] / G2[nz]
|
||
# G=0: zero (charge-neutral cell)
|
||
|
||
# Back to real space: F_n = V̂_H(G_n)/dV → ifftn(F)
|
||
return np.fft.ifftn(V_H_G_full / self.dV).real.ravel()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Hamiltonian actions
|
||
# ------------------------------------------------------------------
|
||
|
||
def potential_G_to_r(self, V_G_full_flat: np.ndarray) -> np.ndarray:
|
||
"""
|
||
IFFT a potential given on the full mesh (N_r,) → real-space (N_r,).
|
||
V_G_full_flat holds f̂(G) values for every G on the mesh.
|
||
"""
|
||
V_G_mesh = V_G_full_flat.reshape(self.mesh)
|
||
return np.fft.ifftn(V_G_mesh / self.dV).real.ravel()
|
||
|
||
def apply_kinetic(self, psi_G: np.ndarray) -> np.ndarray:
|
||
return self.KE_G * psi_G
|
||
|
||
def apply_potential(self, psi_G: np.ndarray, V_r: np.ndarray) -> np.ndarray:
|
||
"""Apply a real-space potential: (Vψ)(G) via FFT convolution."""
|
||
psi_r = self.G_to_r(psi_G)
|
||
return self.r_to_G(V_r * psi_r)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Convenience
|
||
# ------------------------------------------------------------------
|
||
|
||
@staticmethod
|
||
def from_molecule(mol, vacuum: float = 6.0, ecut: float = 20.0):
|
||
"""
|
||
Build a box around the molecule with vacuum padding on each side.
|
||
|
||
Returns (cell, offset) where offset is the shift applied to atomic
|
||
coordinates so the molecule sits at the centre of the box.
|
||
"""
|
||
coords = np.array([a.coords for a in mol.atoms])
|
||
lo = coords.min(axis=0) - vacuum
|
||
hi = coords.max(axis=0) + vacuum
|
||
L = hi - lo
|
||
a_matrix = np.diag(L)
|
||
cell = Cell(a_matrix, ecut)
|
||
offset = -lo # atoms[i].coords + offset → position in cell
|
||
return cell, offset
|
||
|
||
def __repr__(self):
|
||
return (
|
||
f"Cell(volume={self.volume:.2f} Bohr³, ecut={self.ecut:.1f} Ha, "
|
||
f"N_G={self.N_G}, mesh={tuple(int(x) for x in self.mesh)})"
|
||
)
|