Coverage for gpaw/response/site_data.py: 99%
110 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-19 00:19 +0000
1import numpy as np
3from ase.units import Bohr, Hartree
4from ase.neighborlist import natural_cutoffs, build_neighbor_list
6from gpaw.sphere.integrate import (integrate_lebedev,
7 radial_truncation_function,
8 spherical_truncation_function_collection,
9 default_spherical_drcut,
10 find_volume_conserving_lambd)
11from gpaw.response import (ResponseGroundStateAdapter,
12 ResponseGroundStateAdaptable)
13from gpaw.response.localft import (add_spin_polarization,
14 add_LSDA_zeeman_energy)
17class AtomicSites:
18 """Object defining a set of spherical atomic sites."""
20 def __init__(self, indices, radii):
21 """Construct the AtomicSites.
23 Parameters
24 ----------
25 indices : 1D array-like
26 Atomic index A for each site index a.
27 radii : 2D array-like
28 Atomic radius rc for each site index a and partitioning p.
29 """
30 self.A_a = np.asarray(indices)
31 assert self.A_a.ndim == 1
32 assert len(np.unique(self.A_a)) == len(self.A_a)
34 # Parse the input atomic radii
35 rc_ap = np.asarray(radii)
36 assert rc_ap.ndim == 2
37 assert rc_ap.shape[0] == len(self.A_a)
38 # Convert radii to internal units (Å to Bohr)
39 self.rc_ap = rc_ap / Bohr
41 self.npartitions = self.rc_ap.shape[1]
42 self.shape = rc_ap.shape
44 def __len__(self):
45 return len(self.A_a)
48def calculate_site_magnetization(
49 gs: ResponseGroundStateAdaptable,
50 sites: AtomicSites):
51 """Calculate the site magnetization.
53 Returns
54 -------
55 magmom_ap : np.ndarray
56 Magnetic moment in μB of site a under partitioning p, calculated
57 directly from the ground state density.
58 """
59 return AtomicSiteData(gs, sites).calculate_magnetic_moments()
62def calculate_site_zeeman_energy(
63 gs: ResponseGroundStateAdaptable,
64 sites: AtomicSites):
65 """Calculate the site Zeeman energy.
67 Returns
68 -------
69 EZ_ap : np.ndarray
70 Local Zeeman energy in eV of site a under partitioning p, calculated
71 directly from the ground state density.
72 """
73 site_data = AtomicSiteData(gs, sites)
74 return site_data.calculate_zeeman_energies() * Hartree # Ha -> eV
77def get_site_radii_range(gs):
78 """Get the range of valid site radii for the atoms of a given ground state.
80 Returns
81 -------
82 rmin_A : np.ndarray
83 Minimum cutoff radius in Å for each atom A.
84 rmax_A : np.ndarray
85 Maximum cutoff radius in Å for each atom A.
86 """
87 rmin_A, rmax_A = AtomicSiteData.valid_site_radii_range(gs)
88 return rmin_A * Bohr, rmax_A * Bohr # Bohr -> Å
91def maximize_site_magnetization(gs, indices=None):
92 """Find the allowed site radii which maximize the site magnetization.
94 Assumes that m(rc) is maximized for some rc belonging to the interior of
95 the allowed cutoff radii for each atom. Physically, m(rc) has such a
96 maximum only if the spin-polarization of the interstitial region is
97 anti-parallel to the site in its near vicinity.
99 Returns
100 -------
101 rmax_a : np.ndarray
102 Cutoff radius in Å, maximizing the site magnetization for each site a.
103 mmax_a : np.ndarray
104 Site magnetization in μB at its maximum for each site a.
105 """
106 # Calculate the site magnetization as a function of radius
107 rmin_A, rmax_A = get_site_radii_range(gs)
108 if indices is None:
109 indices = range(len(rmin_A))
110 rc_ar = [np.linspace(rmin_A[A], rmax_A[A], 201) for A in indices]
111 magmom_ar = calculate_site_magnetization(gs, AtomicSites(indices, rc_ar))
112 # Maximize the site magnetization
113 rmax_a = np.empty(len(indices), dtype=float)
114 mmax_a = np.empty(len(indices), dtype=float)
115 for a, (rc_r, magmom_r) in enumerate(zip(rc_ar, magmom_ar)):
116 rmax_a[a], mmax_a[a] = maximize(rc_r, magmom_r)
117 return rmax_a, mmax_a
120def maximize(x_x, f_x):
121 """Maximize f(x) on the given interval (returning xmax and f(xmax)).
123 If there is no local maximum on the interior of the interval,
124 we return np.nan.
125 """
126 from gpaw.test import findpeak
127 xmax = f_x.argmax()
128 if xmax == 0 or xmax == len(x_x) - 1:
129 return np.nan, np.nan
130 return findpeak(x_x, f_x)
133class AtomicSiteData:
134 r"""Data object for a set of spherical atomic sites."""
136 def __init__(self, gs: ResponseGroundStateAdaptable, sites: AtomicSites):
137 """Extract atomic site data from a given ground state."""
138 gs = ResponseGroundStateAdapter.from_input(gs)
139 assert self.in_valid_site_radii_range(gs, sites), \
140 'Please provide site radii in the valid range, see '\
141 'gpaw.response.site_data.get_site_radii_range()'
142 self.sites = sites
144 # Extract the scaled positions and micro_setups for each atomic site
145 self.spos_ac = gs.spos_ac[sites.A_a]
146 self.micro_setup_a = [gs.micro_setups[A] for A in sites.A_a]
148 # Extract pseudo density on the fine real-space grid
149 self.finegd = gs.finegd
150 self.nt_sr = gs.nt_sr
152 # Set up the atomic truncation functions which define the sites based
153 # on the coarse real-space grid
154 self.gd = gs.gd
155 self.drcut = default_spherical_drcut(self.gd)
156 self.lambd_ap = np.array(
157 [[find_volume_conserving_lambd(rcut, self.drcut)
158 for rcut in rc_p] for rc_p in sites.rc_ap])
159 self.stfc = spherical_truncation_function_collection(
160 self.finegd, self.spos_ac, sites.rc_ap, self.drcut, self.lambd_ap)
162 @staticmethod
163 def valid_site_radii_range(gs):
164 """For each atom in gs, determine the valid site radii range in Bohr.
166 The lower bound is determined by the spherical truncation width, when
167 truncating integrals on the real-space grid.
168 The upper bound is determined by the distance to the nearest
169 augmentation sphere.
170 """
171 atoms = gs.atoms
172 drcut = default_spherical_drcut(gs.gd)
173 rmin_A = np.array([drcut / 2] * len(atoms))
175 # Find neighbours based on covalent radii
176 cutoffs = natural_cutoffs(atoms, mult=2)
177 neighbourlist = build_neighbor_list(
178 atoms, cutoffs, self_interaction=False, bothways=True)
179 # Determine rmax for each atom
180 augr_A = gs.get_aug_radii()
181 rmax_A = []
182 for A in range(len(atoms)):
183 pos = atoms.positions[A]
184 # Calculate the distance to the augmentation sphere of each
185 # neighbour
186 aug_distances = []
187 for An, offset in zip(*neighbourlist.get_neighbors(A)):
188 posn = atoms.positions[An] + offset @ atoms.get_cell()
189 dist = np.linalg.norm(posn - pos) / Bohr # Å -> Bohr
190 aug_dist = dist - augr_A[An]
191 assert aug_dist > 0.
192 aug_distances.append(aug_dist)
193 # In order for PAW corrections to be valid, we need a sphere of
194 # radius rcut not to overlap with any neighbouring augmentation
195 # spheres
196 rmax_A.append(min(aug_distances))
197 rmax_A = np.array(rmax_A)
199 return rmin_A, rmax_A
201 @staticmethod
202 def in_valid_site_radii_range(gs, sites):
203 rmin_A, rmax_A = AtomicSiteData.valid_site_radii_range(gs)
204 for a, A in enumerate(sites.A_a):
205 if not np.all(
206 np.logical_and(
207 sites.rc_ap[a] > rmin_A[A] - 1e-8,
208 sites.rc_ap[a] < rmax_A[A] + 1e-8)):
209 return False
210 return True
212 def calculate_magnetic_moments(self):
213 """Calculate the magnetic moments at each atomic site."""
214 magmom_ap = self.integrate_local_function(add_spin_polarization)
215 return magmom_ap
217 def calculate_zeeman_energies(self):
218 r"""Calculate the local Zeeman energy E_Z for each atomic site."""
219 EZ_ap = self.integrate_local_function(add_LSDA_zeeman_energy)
220 return EZ_ap
222 def integrate_local_function(self, add_f):
223 r"""Integrate a local function f[n](r) = f(n(r)) over the atomic sites.
225 For every site index a and partitioning p, the integral is defined via
226 a smooth truncation function θ(|r-r_a|<rc_ap):
228 /
229 f_ap = | dr θ(|r-r_a|<rc_ap) f(n(r))
230 /
231 """
232 out_ap = np.zeros(self.sites.shape, dtype=float)
233 self._integrate_pseudo_contribution(add_f, out_ap)
234 self._integrate_paw_correction(add_f, out_ap)
235 return out_ap
237 def _integrate_pseudo_contribution(self, add_f, out_ap):
238 """Calculate the pseudo contribution to the atomic site integrals.
240 For local functions of the density, the pseudo contribution is
241 evaluated by a numerical integration on the real-space grid:
243 ̰ /
244 f_ap = | dr θ(|r-r_a|<rc_ap) f(ñ(r))
245 /
246 """
247 # Evaluate the local function on the real-space grid
248 ft_r = self.finegd.zeros()
249 add_f(self.finegd, self.nt_sr, ft_r)
251 # Integrate θ(|r-r_a|<rc_ap) f(ñ(r))
252 ftdict_ap = {a: np.empty(self.sites.npartitions)
253 for a in range(len(self.sites))}
254 self.stfc.integrate(ft_r, ftdict_ap)
256 # Add pseudo contribution to the output array
257 for a in range(len(self.sites)):
258 out_ap[a] += ftdict_ap[a]
260 def _integrate_paw_correction(self, add_f, out_ap):
261 """Calculate the PAW correction to an atomic site integral.
263 The PAW correction is evaluated on the atom centered radial grid, using
264 the all-electron and pseudo densities generated from the partial waves:
266 /
267 Δf_ap = | r^2 dr θ(r<rc_ap) [f(n_a(r)) - f(ñ_a(r))]
268 /
269 """
270 for a, (micro_setup, rc_p, lambd_p) in enumerate(zip(
271 self.micro_setup_a, self.sites.rc_ap, self.lambd_ap)):
272 # Evaluate the PAW correction and integrate angular components
273 df_ng = micro_setup.evaluate_paw_correction(add_f)
274 df_g = integrate_lebedev(df_ng)
275 for p, (rcut, lambd) in enumerate(zip(rc_p, lambd_p)):
276 # Evaluate the smooth truncation function
277 theta_g = radial_truncation_function(
278 micro_setup.rgd.r_g, rcut, self.drcut, lambd)
279 # Integrate θ(r) Δf(r) on the radial grid
280 out_ap[a, p] += micro_setup.rgd.integrate_trapz(df_g * theta_g)