# coding: utf-8
"""This module contains the class describing densities in real space on uniform 3D meshes."""
import numpy as np
import collections
import pymatgen.core.units as pmgu
import os
from collections import OrderedDict
from monty.collections import AttrDict
from monty.functools import lazy_property
from monty.string import is_string, marquee
from monty.termcolor import cprint
from monty.inspect import all_subclasses
from pymatgen.io.vasp.inputs import Poscar
from pymatgen.io.vasp.outputs import Chgcar
from pymatgen.core.units import bohr_to_angstrom
from abipy.core.structure import Structure
from abipy.core.mesh3d import Mesh3D
from abipy.core.func1d import Function1D
from abipy.core.mixins import Has_Structure
from abipy.tools import duck
from abipy.tools.numtools import transpose_last3dims
from abipy.tools.plotting import add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt
from abipy.iotools import Visualizer, xsf, ETSF_Reader, cube
__all__ = [
"Density",
"VxcPotential",
"VhartreePotential",
"VhxcPotential",
"VksPotential",
]
def latexlabel_ispden(ispden, nspden):
"""Return latexlabel from ``ispden`` with given ``nspden``."""
if nspden == 1:
return None
elif nspden == 2:
return {k: v.replace("myuparrow", "uparrow") for k, v in
{0: r"$\sigma=\myuparrow$", 1: r"$\sigma=\downarrow$"}.items()}[ispden]
else:
raise NotImplementedError()
class _Field(Has_Structure):
"""
Base class representing a set of spin-dependent scalar fields generated by electrons (e.g. densities, potentials).
Data is represented on a homogenous real-space mesh.
A ``_Field`` has a structure object and provides helper functions to perform common operations
such computing integrals, FFT transforms ...
The following attributes must be defined by the subclass:
netcdf_name: String with the name of the netcdf variable.
latex_label: String used in plot to set the axis label.
"""
netcdf_name = "Unknown"
latex_label = " "
@classmethod
def from_file(cls, filepath):
"""Initialize the object from a netCDF file."""
with FieldReader(filepath) as r:
return r.read_denpot(varname=cls.netcdf_name, field_cls=cls)
def __init__(self, nspinor, nsppol, nspden, datar, structure, iorder="c"):
"""
Args:
nspinor: Number of spinorial components.
nsppol: Number of spins.
nspden: Number of spin density components.
datar: [nspden, nx, ny, nz] array with the scalar field in real space.
See also ``read_denpot``.
structure: |Structure| object describing the crystalline structure.
iorder: Order of the array. "c" for C ordering, "f" for Fortran ordering.
"""
self.nspinor, self.nsppol, self.nspden = nspinor, nsppol, nspden
# Convert to Abipy Structure.
self._structure = Structure.as_structure(structure)
iorder = iorder.lower()
assert iorder in ["f", "c"]
if iorder == "f":
# (z,y,x) --> (x,y,z)
datar = transpose_last3dims(datar)
# Init Mesh3D
mesh_shape = datar.shape[-3:]
self._mesh = Mesh3D(mesh_shape, structure.lattice.matrix)
# Make sure we have the correct shape.
self._datar = np.reshape(datar, (nspden,) + self.mesh.shape)
def __len__(self):
return len(self.datar)
def __str__(self):
return self.to_string()
def _check_and_get_datar(self, other):
"""Perform consistency check and return datar values."""
if not isinstance(other, _Field):
try:
return self.__class__, float(other)
except Exception:
raise TypeError('object of class %s is not an instance of _Field and cannot be converted to float' %
(other.__class__))
if any([self.nspinor != other.nspinor, self.nsppol != other.nsppol, self.nspden != other.nspden,
self.structure != other.structure, self.mesh != other.mesh]):
raise ValueError('Incompatible scalar fields')
new_cls = self.__class__ if isinstance(other, self.__class__) else _Field
return new_cls, other.datar
def __add__(self, other):
"""self + other"""
new_cls, datar = self._check_and_get_datar(other)
return new_cls(nspinor=self.nspinor, nsppol=self.nsppol, nspden=self.nspden,
datar=self.datar + datar,
structure=self.structure, iorder="c")
def __sub__(self, other):
"""self - other"""
new_cls, datar = self._check_and_get_datar(other)
return new_cls(nspinor=self.nspinor, nsppol=self.nsppol, nspden=self.nspden,
datar=self.datar - datar,
structure=self.structure, iorder="c")
def __mul__(self, other):
"""self * other"""
new_cls, datar = self._check_and_get_datar(other)
return new_cls(nspinor=self.nspinor, nsppol=self.nsppol, nspden=self.nspden,
datar=self.datar * datar,
structure=self.structure, iorder="c")
__rmul__ = __mul__
def __truediv__(self, other):
"""self / other"""
new_cls, datar = self._check_and_get_datar(other)
return new_cls(nspinor=self.nspinor, nsppol=self.nsppol, nspden=self.nspden,
datar=self.datar / datar,
structure=self.structure, iorder="c")
__div__ = __truediv__
def __neg__(self):
"""-self"""
return self.__class__(nspinor=self.nspinor, nsppol=self.nsppol, nspden=self.nspden,
datar=-self.datar,
structure=self.structure, iorder="c")
def __abs__(self):
"""abs(self)"""
return self.__class__(nspinor=self.nspinor, nsppol=self.nsppol, nspden=self.nspden,
datar=np.abs(self.datar),
structure=self.structure, iorder="c")
@property
def structure(self):
"""|Structure| object."""
return self._structure
def to_string(self, verbose=0, title=None):
"""String representation"""
lines = []; app = lines.append
if title is not None: app(marquee(title), mark="=")
app("%s: nspinor: %i, nsppol: %i, nspden: %i" %
(self.__class__.__name__, self.nspinor, self.nsppol, self.nspden))
app(self.mesh.to_string(verbose=verbose))
if verbose > 0:
app(self.structure.to_string(verbose=verbose))
return "\n".join(lines)
@property
def datar(self):
"""|numpy-array| with data in real space. shape: [nspden, nx, ny, nz]"""
return self._datar
@lazy_property
def datag(self):
"""|numpy-array| with data in reciprocal space. shape: [nspden, nx, ny, nz]"""
# FFT R --> G.
return self.mesh.fft_r2g(self.datar)
@property
def mesh(self):
"""|Mesh3D| object. datar and datag are defined on this mesh."""
return self._mesh
@property
def shape(self):
"""Shape of the array."""
assert self.datar.shape == self.datag.shape
return self.datar.shape
@property
def nx(self):
"""Number of points along x."""
return self.mesh.nx
@property
def ny(self):
"""Number of points along y."""
return self.mesh.ny
@property
def nz(self):
"""Number of points along z."""
return self.mesh.nz
@property
def is_density_like(self):
"""True if field is density-like."""
return isinstance(self, _DensityField)
@property
def is_potential_like(self):
"""True if field is potential-like."""
return isinstance(self, _PotentialField)
@property
def is_collinear(self):
"""True if collinear i.e. nspinor==1."""
return self.nspinor == 1
@staticmethod
def _check_space(space):
"""Helper function used in __add__ ... methods to check Consistency."""
space = space.lower()
if space not in ("r", "g"):
raise ValueError("Wrong space %s" % space)
return space
def mean(self, space="r", axis=0):
"""
Returns the average of the array elements along the given axis.
"""
if "r" == self._check_space(space):
return self.datar.mean(axis=axis)
else:
return self.datag.mean(axis=axis)
def std(self, space="r", axis=0):
"""
Returns the standard deviation of the array elements along the given axis.
"""
if "r" == self._check_space(space):
return self.datar.std(axis=axis)
else:
return self.datag.std(axis=axis)
def export(self, filename, visu=None, verbose=1):
"""
Export the real space data to file filename.
Args:
filename: String specifying the file path and the file format.
The format is defined by the file extension. filename="prefix.xsf", for example,
will produce a file in XSF format (xcrysden_).
An *empty* prefix, e.g. ".xsf" makes the code use a temporary file.
visu:
:class:`Visualizer` subclass. By default, this method returns the first available
visualizer that supports the given file format. If visu is not None, an
instance of visu is returned. See :class:`Visualizer` for the list of
applications and formats supported.
verbose: Verbosity level
Returns:
Instance of :class:`Visualizer`
"""
if "." not in filename:
raise ValueError("Cannot detect file extension in filename: %s " % filename)
tokens = filename.strip().split(".")
ext = tokens[-1]
if verbose:
print("tokens", tokens, "ext", ext)
if not tokens[0]:
# filename == ".ext" ==> Create temporary file.
# dir = os.getcwd() is needed when we invoke the method from a notebook.
# nbworkdir in cwd is needed when we invoke the method from a notebook.
from abipy.core.globals import abinb_mkstemp
_, filename = abinb_mkstemp(suffix="." + ext, text=True)
with open(filename, mode="wt") as fh:
if ext == "xsf":
# xcrysden
xsf.xsf_write_structure(fh, self.structure)
xsf.xsf_write_data(fh, self.structure, self.datar, add_replicas=True)
#elif ext == "POSCAR":
else:
raise NotImplementedError("extension %s is not supported." % ext)
if visu is None:
return Visualizer.from_file(filename)
else:
return visu(filename)
def visualize(self, appname):
"""
Visualize data with visualizer.
See :class:`Visualizer` for the list of applications and formats supported.
"""
visu = Visualizer.from_name(appname)
# Try to export data to one of the formats supported by the visualizer
# Use a temporary file (note "." + ext)
for ext in visu.supported_extensions():
ext = "." + ext
try:
return self.export(ext, visu=visu)()
except visu.Error:
pass
else:
raise visu.Error("Don't know how to export data for visualizer %s" % appname)
def get_interpolator(self):
"""
Return an interpolator object that interpolates periodic functions in real space.
"""
from abipy.tools.numtools import BlochRegularGridInterpolator
return BlochRegularGridInterpolator(self.structure, self.datar)
#def fourier_interp(self, new_mesh):
#intp_datar = self.mesh.fourier_interp(self.datar, new_mesh, inspace="r")
#return self.__class__(self.nspinor, self.nsppol, self.nspden, self.structure, intp_datar)
#def braket_waves(self, bra_wave, ket_wave):
# """
# Compute the matrix element of <bra_wave| self.datar |ket_wave> in real space.
# """
# bra_ur = bra_wave.get_ur_mesh(self.mesh, copy=False)
# ket_ur = ket_wave.get_ur_mesh(self.mesh, copy=False)
# if self.nspinor == 1:
# assert bra_wave.spin == ket_wave.spin
# v = self.mesh.integrate(bra_ur.conj() * self.datar[bra_wave.spin] * ket_ur)
# return v / self.structure.volume
# else:
# raise NotImplementedError("nspinor != 1 not implenented")
@add_fig_kwargs
def plot_line(self, point1, point2, num=200, cartesian=False, ax=None, fontsize=12, **kwargs):
"""
Plot (interpolated) density/potential in real space along a line defined by ``point1`` and ``point2``.
Args:
point1: First point of the line. Accepts 3d vector or integer.
The vector is in reduced coordinates unless ``cartesian`` is True.
If integer, the first point of the line is given by the i-th site of the structure
e.g. ``point1=0, point2=1`` gives the line passing through the first two atoms.
point2: Second point of the line. Same API as ``point1``.
num: Number of points sampled along the line.
cartesian: By default, ``point1`` and ``point1`` are interpreted as points in fractional
coordinates (if not integers). Use True to pass points in cartesian coordinates.
ax: |matplotlib-Axes| or None if a new figure should be created.
fontsize: legend and title fontsize.
Return: |matplotlib-Figure|
"""
# Interpolate along line.
r = self.get_interpolator().eval_line(point1, point2, num=num, cartesian=cartesian)
# Plot data.
ax, fig, plt = get_ax_fig_plt(ax=ax)
for ispden in range(self.nspden):
ax.plot(r.dist, r.values[ispden], label=latexlabel_ispden(ispden, self.nspden))
ax.grid(True)
if r.site1 is None:
ax.set_xlabel("Distance from point %s [Angstrom]" % str(point1))
else:
cs = "[%+.5f, %+.5f, %+.5f]" % tuple(r.site1.frac_coords)
ax.set_xlabel("Distance in Angstrom from %s at %s" % (r.site1.specie.symbol, cs))
ax.set_ylabel(self.latex_label)
if self.nspden > 1:
ax.legend(loc="best", fontsize=fontsize, shadow=True)
return fig
@add_fig_kwargs
def plot_line_neighbors(self, site_index, radius, num=200, max_nn=10, fontsize=12, **kwargs):
"""
Plot (interpolated) density/potential in real space along the lines connecting
an atom specified by ``site_index`` and all neighbors within a sphere of given ``radius``.
.. warning::
This routine can produce lots of plots!
Be careful with the value of ``radius``. See also ``max_nn``.
Args:
site_index: Index of the atom in the structure.
radius: Radius of the sphere in Angstrom
num: Number of points sampled along the line.
max_nn: By default, only the first `max_nn` neighbors are showed.
fontsize: legend and title fontsize
Return: |matplotlib-Figure|
"""
site = self.structure[site_index]
nn_list = self.structure.get_neighbors_old(site, radius, include_index=True)
if not nn_list:
cprint("Zero neighbors found for radius %s Ang. Returning None." % radius, "yellow")
return None
# Sorte sites by distance.
nn_list = list(sorted(nn_list, key=lambda t: t.nn_distance))
if max_nn is not None and len(nn_list) > max_nn:
cprint("For radius %s, found %s neighbors but only max_nn %s sites are show." %
(radius, len(nn_list), max_nn), "yellow")
nn_list = nn_list[:max_nn]
# Get grid of axes.
nrows, ncols = len(nn_list), 1
ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
sharex=True, sharey=True, squeeze=True)
ax_list = ax_list.ravel()
interpolator = self.get_interpolator()
for i, (nn, ax) in enumerate(zip(nn_list, ax_list)):
#nn_site, nn_dist, nn_sc_index = nn
nn_site = nn
nn_dist = nn.nn_distance
nn_sc_index = nn.index
title = "%s, %s, dist=%.3f A" % (nn_site.species_string, str(nn_site.frac_coords), nn_dist)
r = interpolator.eval_line(site.frac_coords, nn_site.frac_coords, num=num, kpoint=None)
for ispden in range(self.nspden):
ax.plot(r.dist, r.values[ispden],
label=latexlabel_ispden(ispden, self.nspden) if i == 0 else None)
ax.set_title(title, fontsize=fontsize)
ax.grid(True)
if i == nrows - 1:
ax.set_xlabel("Distance from site_index %s [Angstrom]" % site_index)
ax.set_ylabel(self.latex_label)
if self.nspden > 1:
ax.legend(loc="best", fontsize=fontsize, shadow=True)
return fig
def integrate_in_spheres(self, rcut_symbol=None, out=False):
"""
Integrate field (e.g. density/potential) inside atom-centered spheres of given radius.
Can be used to get a rough estimate of the charge/magnetization associated to a given site.
Args:
rcut_symbol: dictionary mapping chemical element to the radius of the sphere in Angstrom.
or number if each element should have the same sphere. If None, covalent radii are used.
out: Set it to False to disable output of final results
Return:
|pandas-DataFrame| with computed results (integrated density, integrated magnetization, ...)
"""
# Initialize rcut_symbol map.
if rcut_symbol is None:
from pymatgen.analysis.molecule_structure_comparator import CovalentRadius
rcut_symbol = {s: CovalentRadius.radius[s] for s in self.structure.symbol_set}
#rcut_symbol = {s: 2 for s in self.structure.symbol_set}
#rcut_symbol = {s: 1 for s in self.structure.symbol_set}
elif duck.is_number_like(rcut_symbol):
rcut_symbol = {s: float(rcut_symbol) for s in self.structure.symbol_set}
# Spline bessel integrals.
datag = np.reshape(self.datag, (self.nspden, -1))
#print("datag[0]", datag[0, 0] * self.structure.volume, datag.shape)
gvecs = self.mesh.gvecs
gmods = self.mesh.gmods
gmax = gmods.max()
from abipy.tools import bessel
splines = {s: bessel.spline_int_jlqr(0, gmax, rcut_symbol[s]) for s in self.structure.symbol_set}
# 4 pi sum_G n(G) e^{iGRo} int_0^{rcut} r**2 j_l(Gr} dr
rows = []
for iatom, site in enumerate(self.structure):
symbol = site.specie.symbol
phases = np.exp(2j * np.pi * np.dot(gvecs, site.frac_coords))
fg = datag * phases * splines[symbol](gmods)
res_nspden = np.sum(fg, axis=1) * (4 * np.pi)
#print("result:", res_nspden, res_nspden.shape, (datag * fg).shape)
# Compute densities and magnetization.
ntot, nup, ndown, mx, my, mz = 6 * (None,)
if self.nspinor == 1:
res_nspden = res_nspden.real
if self.nspden == 1:
ntot = res_nspden[0]
elif self.nspden == 2:
nup, ndown = res_nspden
ntot, mz = nup + ndown, nup - ndown
elif self.nspinor == 2:
raise NotImplementedError()
#ntot, mx, my, mz = scalvec_from_spinmat(res_nspden)
nup, ndown = 0.5 * (ntot + mz), 0.5 * (ntot - mz)
# Fill DataFrame row.
rows.append(OrderedDict([
("iatom", iatom), ("symbol", symbol),
("ntot", ntot), ("nup", nup), ("ndown", ndown),
("mx", mx), ("my", my), ("mz", mz),
("rsph_ang", rcut_symbol[symbol]), ("frac_coords", site.frac_coords),
]))
import pandas as pd
df = pd.DataFrame(rows, columns=list(rows[0].keys()))
# Use iatom as index and remove columns with None.
df = df.set_index("iatom").dropna(axis="columns", how='any')
if out:
print(self.structure)
print()
print(df)
return df
class _DensityField(_Field):
"""Base class for density-like fields."""
def core_density_from_file(filepath):
"""
Parses a file and extracts two numpy array, containing radii and the core densities, respectively.
Can be used to provide the core densities to Density.ae_core_density_on_mesh
Supported file types: fc, rhoc
"""
ext = os.path.splitext(filepath)[1]
if ext == '.fc':
with open(filepath, 'r') as f:
lines = f.readlines()
r, rho = [], []
for l in lines[1:]:
if not l.strip():
continue
if l.startswith('<JSON>'):
break
l = l.split()
r.append(float(l[0]))
rho.append(float(l[1]))
return np.array(r) * bohr_to_angstrom, np.array(rho) / (4.0*np.pi) / (bohr_to_angstrom ** 3)
elif ext == '.rhoc':
rhoc = np.loadtxt(filepath)
return rhoc[:, 0] * bohr_to_angstrom, rhoc[:,1] / (4.0*np.pi) / (bohr_to_angstrom ** 3)
else:
raise ValueError('Exension not supported: {}'.format(ext))
[docs]class Density(_DensityField):
"""
Electronic density.
.. note::
Unlike in Abinit, datar[nspden] contains the up/down components if nsppol = 2
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: Density
"""
netcdf_name = "density"
latex_label = "Density [$e/A^3$]"
[docs] @classmethod
def ae_core_density_on_mesh(cls, valence_density, structure, rhoc, maxr=2.0, nelec=None, tol=0.01,
method='get_sites_in_sphere', small_dist_mesh=(8, 8, 8), small_dist_factor=1.5):
"""
Initialize the all electron core density of the structure from the pseudopotentials *rhoc* files.
For points close to the atoms, the value at the grid point would be defined as the average on a finer grid
in the neghibourhood of the point.
Args:
valence_density: a Density object representing the valence charge density
structure: the structure for which the total core density will be calculated
rhoc: *rhoc* files for the elements in structure. Can be a list with len=len(structure) or a
dictionary {element: rhoc}. Distances should be in Angstrom and densities in e/A^3.
maxr: maximum radius for the distance between grid points and atoms.
nelec: if given a check will be perfomed to verify that the integrated density will be within 1% error
with respect to nelec. The total density will also be rescaled to fit exactly the given number.
tol: tolerance above which the system will raise an exception if the integrated density doesn't sum up
to the value specified in nelec. Default 0.01 (1% error).
method: different methods to perform the calculation:
* get_sites_in_sphere: based on ``Structure.get_sites_in_sphere``.
* mesh3d_dist_gridpoints: based on ``Mesh3D.dist_gridpoints_in_spheres``. Generally faster than
``get_sites_in_sphere``, but tests can be made for specific cases.
* get_sites_in_sphere_legacy: as get_sites_in_sphere, but part of the procedure is not vectorized
* mesh3d_dist_gridpoints_legacy: as mesh3d_dist_gridpoints, but part of the procedure is not vectorized
small_dist_mesh: defines the finer mesh.
small_dist_factor: defines the maximum distance for which a point should be considered close enough to
switch to the finer grid method. Note that this is a factor, the distance is defined with respect
to the size of the cell.
"""
rhoc_atom_splines = [None]*len(structure)
if isinstance(rhoc, (list, tuple)):
if len(structure) != len(rhoc):
raise ValueError('Number of rhoc files should be equal to the number of sites in the structure')
elif isinstance(rhoc, collections.abc.Mapping):
atoms_symbols = [elmt.symbol for elmt in structure.composition]
if not np.all([atom in rhoc for atom in atoms_symbols]):
raise ValueError('The rhoc files should be provided for all the atoms in the structure')
rhoc = [rhoc[site.specie.symbol] for site in structure]
else:
raise ValueError('Unsuported format for rhoc')
for ir, r in enumerate(rhoc):
func1d = Function1D(r[0], r[1])
rhoc_atom_splines[ir] = func1d.spline
# if maxr is negative find the minimum radius so that for all elements the density is zero.
if maxr < 0:
abs_max = abs(maxr)
for r in rhoc:
try:
ind = np.min(np.where(r[1] == 0))
except Exception:
ind = -1
if r[0][ind] > maxr:
maxr = r[0][ind]
if maxr > abs_max:
maxr = abs_max
break
core_den = np.zeros_like(valence_density.datar)
dvx = valence_density.mesh.dvx
dvy = valence_density.mesh.dvy
dvz = valence_density.mesh.dvz
maxdiag = max([np.linalg.norm(dvx+dvy+dvz),
np.linalg.norm(dvx+dvy-dvz),
np.linalg.norm(dvx-dvy+dvz),
np.linalg.norm(dvx-dvy-dvz)])
smallradius = small_dist_factor*maxdiag
# The vectorized methods are faster. Keep the older methods for cross checks of the implementation for the
# time being
if method == 'get_sites_in_sphere_legacy':
for ix in range(valence_density.mesh.nx):
for iy in range(valence_density.mesh.ny):
for iz in range(valence_density.mesh.nz):
rpoint = valence_density.mesh.rpoint(ix=ix, iy=iy, iz=iz)
# TODO: optimize this !
sites = structure.get_sites_in_sphere(pt=rpoint, r=maxr, include_index=True)
for site, dist, site_index in sites:
if dist > smallradius:
core_den[0, ix, iy, iz] += rhoc_atom_splines[site_index](dist)
# For small distances, integrate over the small volume dv around the point as the core
# density is extremely high close to the atom
else:
total = 0.0
nnx, nny, nnz = small_dist_mesh
ddvx = dvx/nnx
ddvy = dvy/nny
ddvz = dvz/nnz
rpi = rpoint - 0.5 * (dvx + dvy + dvz) + 0.5*ddvx + 0.5*ddvy + 0.5*ddvz
for iix in range(nnx):
for iiy in range(nny):
for iiz in range(nnz):
rpoint2 = rpi + iix*ddvx + iiy*ddvy + iiz*ddvz
dist2 = np.linalg.norm(rpoint2 - site.coords)
total += rhoc_atom_splines[site_index](dist2)
total /= (nnx*nny*nnz)
core_den[0, ix, iy, iz] += total
elif method == 'mesh3d_dist_gridpoints_legacy':
site_coords = [site.coords for site in structure]
dist_gridpoints_sites = valence_density.mesh.dist_gridpoints_in_spheres(points=site_coords, radius=maxr)
for isite, dist_gridpoints_site in enumerate(dist_gridpoints_sites):
for igp_uc, dist, igp in dist_gridpoints_site:
if dist > smallradius:
core_den[0, igp_uc[0], igp_uc[1], igp_uc[2]] += rhoc_atom_splines[isite](dist)
# For small distances, integrate over the small volume dv around the point as the core density
# is extremely high close to the atom
else:
total = 0.0
nnx, nny, nnz = small_dist_mesh
ddvx = dvx/nnx
ddvy = dvy/nny
ddvz = dvz/nnz
rpoint = valence_density.mesh.rpoint(ix=igp[0], iy=igp[1], iz=igp[2])
rpi = rpoint - 0.5 * (dvx + dvy + dvz) + 0.5*ddvx + 0.5*ddvy + 0.5*ddvz
for iix in range(nnx):
for iiy in range(nny):
for iiz in range(nnz):
rpoint2 = rpi + iix*ddvx + iiy*ddvy + iiz*ddvz
dist2 = np.linalg.norm(rpoint2 - site_coords[isite])
total += rhoc_atom_splines[isite](dist2)
total /= (nnx*nny*nnz)
core_den[0, igp_uc[0], igp_uc[1], igp_uc[2]] += total
elif method == 'mesh3d_dist_gridpoints':
import time
site_coords = [site.coords for site in structure]
start = time.time()
dist_gridpoints_sites = valence_density.mesh.dist_gridpoints_in_spheres(points=site_coords, radius=maxr)
nnx, nny, nnz = small_dist_mesh
meshgrid = np.meshgrid(np.linspace(-0.5, 0.5, nnx, endpoint=False) + 0.5 / nnx,
np.linspace(-0.5, 0.5, nny, endpoint=False) + 0.5 / nny,
np.linspace(-0.5, 0.5, nnz, endpoint=False) + 0.5 / nnz)
coords_grid = np.outer(meshgrid[0], dvx) + np.outer(meshgrid[1], dvy) + np.outer(meshgrid[2], dvz)
for isite, dist_gridpoints_site in enumerate(dist_gridpoints_sites):
for igp_uc, dist, igp in dist_gridpoints_site:
if dist > smallradius:
core_den[0, igp_uc[0], igp_uc[1], igp_uc[2]] += rhoc_atom_splines[isite](dist)
# For small distances, integrate over the small volume dv around the point as the core density
# is extremely high close to the atom
else:
total = 0.0
rpoint = valence_density.mesh.rpoint(ix=igp[0], iy=igp[1], iz=igp[2])
grid_loc = rpoint + coords_grid
distances = np.linalg.norm(grid_loc-site_coords[isite], axis=1)
total = np.sum(rhoc_atom_splines[isite](distances))
total /= (nnx*nny*nnz)
core_den[0, igp_uc[0], igp_uc[1], igp_uc[2]] += total
elif method == 'get_sites_in_sphere':
nnx, nny, nnz = small_dist_mesh
meshgrid = np.meshgrid(np.linspace(-0.5, 0.5, nnx, endpoint=False) + 0.5 / nnx,
np.linspace(-0.5, 0.5, nny, endpoint=False) + 0.5 / nny,
np.linspace(-0.5, 0.5, nnz, endpoint=False) + 0.5 / nnz)
coords_grid = np.outer(meshgrid[0], dvx) + np.outer(meshgrid[1], dvy) + np.outer(meshgrid[2], dvz)
for ix in range(valence_density.mesh.nx):
for iy in range(valence_density.mesh.ny):
for iz in range(valence_density.mesh.nz):
rpoint = valence_density.mesh.rpoint(ix=ix, iy=iy, iz=iz)
sites = structure.get_sites_in_sphere(pt=rpoint, r=maxr, include_index=True)
#for site, dist, site_index in sites:
for site in sites:
dist, site_index = site.nn_distance, site.index
if dist > smallradius:
core_den[0, ix, iy, iz] += rhoc_atom_splines[site_index](dist)
# For small distances, integrate over the small volume dv around the point as the core
# density is extremely high close to the atom
else:
total = 0.0
grid_loc = rpoint + coords_grid
distances = np.linalg.norm(grid_loc - site.coords, axis=1)
total = np.sum(rhoc_atom_splines[site_index](distances))
total /= (nnx * nny * nnz)
core_den[0, ix, iy, iz] += total
else:
raise ValueError('Method "{}" is not allowed'.format(method))
if nelec is not None:
sum_elec = np.sum(core_den) * valence_density.mesh.dv
diff = np.abs(sum_elec-nelec) / nelec
if diff > tol:
raise ValueError('Summed electrons is different from the actual number of electrons by '
'more than {:.2f}% : {:.2f}%'.format(tol*100, diff*100))
core_den = core_den / sum_elec * nelec
return cls(nspinor=valence_density.nspinor, nsppol=valence_density.nsppol, nspden=valence_density.nspden,
datar=core_den, structure=structure, iorder='c')
[docs] def get_nelect(self, spin=None):
"""
Returns the number of electrons with given spin.
If spin is None, the total number of electrons is computed.
"""
if self.is_collinear:
nelect = self.mesh.integrate(self.datar)
return np.sum(nelect) if spin is None else nelect[spin]
else:
raise NotImplementedError()
return self.mesh.integrate(self.datar[0])
[docs] @lazy_property
def total_rhor(self):
"""
|numpy-array| with the total density in real space on the FFT mesh.
"""
if self.is_collinear:
if self.nsppol == 1:
if self.nspden == 2: raise NotImplementedError()
return self.datar[0]
elif self.nsppol == 2:
return self.datar[0] + self.datar[1]
else:
raise ValueError("You should not be here")
raise NotImplementedError("Non collinear case.")
[docs] def total_rhor_as_density(self):
"""Return a :class:`Density` object with the total density."""
return self.__class__(nspinor=1, nsppol=1, nspden=1, datar=self.total_rhor,
structure=self.structure, iorder="c")
[docs] @lazy_property
def total_rhog(self):
"""|numpy-array| with the total density in G-space."""
# FFT R --> G.
return self.mesh.fft_r2g(self.total_rhor)
[docs] @lazy_property
def magnetization_field(self):
"""
|numpy-array| with the magnetization field in real space on the FFT mesh:
* 0 if spin-unpolarized calculation
* spin_up - spin_down if collinear spin-polarized
* numpy array with (mx, my, mz) components if non-collinear magnetism
"""
if self.is_collinear:
if self.nsppol == 1 and self.nspden == 1:
# zero magnetization by definition.
return self.mesh.zeros()
else:
# spin_up - spin_down.
return self.datar[0] - self.datar[1]
else:
# mx, my, mz
return self.datar[1:]
[docs] @lazy_property
def magnetization(self):
"""
Magnetization field integrated over the unit cell.
Scalar if collinear, vector with (mx, my, mz) components if non-collinear.
"""
return self.mesh.integrate(self.magnetization_field)
[docs] @lazy_property
def nelect_updown(self):
"""
Tuple with the number of electrons in the up/down channel.
Return (None, None) if non-collinear.
"""
if not self.is_collinear:
return None, None
if self.nsppol == 1:
if self.nspden == 2: raise NotImplementedError()
nup = ndown = self.mesh.integrate(self.datar[0]/2)
else:
nup = self.mesh.integrate(self.datar[0])
ndown = self.mesh.integrate(self.datar[1])
return nup, ndown
[docs] @lazy_property
def zeta(self):
"""
|numpy-array| with Magnetization(r) / total_density(r)
"""
return self.magnetization * np.where(self.total_rhor > 1e-16, 1 / self.total_rhor, 0.0)
#def vhartree(self):
# """
# Solve the Poisson's equation in reciprocal space.
# returns:
# (vhr, vhg) Hartree potential in real, reciprocal space.
# """
# # Compute |G| for each G in the mesh and treat G=0.
# gvecs = self.mesh.gvecs
# gwork = self.mesh.zeros().ravel()
# gnorm = self.structure.gnorm(gvec)
# for idx, gg in enumerate(gvecs):
# #gnorm = self.structure.gnorm(gg)
# gnorm = 1.0 # self.structure.gnorm(gg)
# #gg = np.atleast_2d(gg)
# #mv = np.dot(self.structure.gmet, gg.T)
# #norm2 = 2*np.pi * np.dot(gg, mv)
# #gnorm = np.sqrt(norm2)
# #print gg, gnorm
# if idx != 0:
# gwork[idx] = 4*np.pi/gnorm
# else:
# gwork[idx] = 0.0
# new_shape = self.mesh.ndivs
# gwork = np.reshape(gwork, new_shape)
# #gwork = self.mesh.reshape(gwork)
# # FFT to obtain vh in real space.
# vhg = self.total_rhog * gwork
# vhr = self.mesh.fft_g2r(vhg, fg_ishifted=False)
# return vhr, vhg
[docs] def export_to_cube(self, filename, spin='total'):
"""
Export real space density to CUBE file ``filename``.
"""
if spin != 'total':
raise ValueError('Argument "spin" should be "total"')
with open(filename, mode="wt") as fh:
cube.cube_write_structure_mesh(file=fh, structure=self.structure, mesh=self.mesh)
cube.cube_write_data(file=fh, data=self.total_rhor, mesh=self.mesh)
[docs] @classmethod
def from_cube(cls, filename, spin='total'):
"""
Read real space density to CUBE file ``filename``.
Return new :class:`Density` instance.
"""
if spin != 'total':
raise ValueError('Argument "spin" should be "total"')
structure, mesh, datar = cube.cube_read_structure_mesh_data(file=filename)
return cls(nspinor=1, nsppol=1, nspden=1, datar=datar, structure=structure, iorder="c")
#@lazy_property
#def kinden(self):
#"""Compute the kinetic energy density in real- and reciprocal-space."""
#return kindr, kindgg
[docs] def to_chgcar(self, filename=None):
"""
Convert a :class:`Density` object into a ``Chgar`` object.
If ``filename`` is not None, density is written to this file in Chgar format
Return:
:class:`Chgcar` instance.
.. note::
From: http://cms.mpi.univie.ac.at/vasp/vasp/CHGCAR_file.html:
This file contains the total charge density multiplied by the volume
For spinpolarized calculations, two sets of data can be found in the CHGCAR file.
The first set contains the total charge density (spin up plus spin down),
the second one the magnetization density (spin up minus spin down).
For non collinear calculations the CHGCAR file contains the total charge density
and the magnetisation density in the x, y and z direction in this order.
"""
myrhor = self.datar * self.structure.volume
if self.nspinor == 1:
if self.nsppol == 1:
data_dict = {"total": myrhor[0]}
if self.nsppol == 2:
data_dict = {"total": myrhor[0] + myrhor[1], "diff": myrhor[0] - myrhor[1]}
elif self.nspinor == 2:
raise NotImplementedError("pymatgen Chgcar does not implement nspinor == 2")
chgcar = Chgcar(Poscar(self.structure), data_dict)
if filename is not None:
chgcar.write_file(filename)
return chgcar
[docs] @classmethod
def from_chgcar_poscar(cls, chgcar, poscar):
"""
Build a :class`Density` object from Vasp data.
Args:
chgcar: Either string with the name of a CHGCAR file or :class:`Chgcar` pymatgen object.
poscar: Either string with the name of a POSCAR file or :class:`Poscar` pymatgen object.
.. warning:
The present version does not support non-collinear calculations.
The Chgcar object provided by pymatgen does not provided enough information
to understand if the calculation is collinear or no.
"""
if is_string(chgcar):
chgcar = Chgcar.from_file(chgcar)
if is_string(poscar):
poscar = Poscar.from_file(poscar, check_for_POTCAR=False, read_velocities=False)
nx, ny, nz = chgcar.dim
nspinor = 1
nsppol = 2 if chgcar.is_spin_polarized else 1
nspden = 2 if nsppol == 2 else 1
# Convert pymatgen chgcar data --> abipy representation.
abipy_datar = np.empty((nspden, nx, ny, nz))
if nspinor == 1:
if nsppol == 1:
abipy_datar = chgcar.data["total"]
elif nsppol == 2:
total, diff = chgcar.data["total"], chgcar.data["diff"]
abipy_datar[0] = 0.5 * (total + diff)
abipy_datar[1] = 0.5 * (total - diff)
else:
raise ValueError("Wrong nspden %s" % nspden)
else:
raise NotImplementedError("nspinor == 2 requires more info in Chgcar")
# density in Chgcar is multiplied by volume!
abipy_datar /= poscar.structure.volume
return cls(nspinor=nspinor, nsppol=nsppol, nspden=nspden, datar=abipy_datar,
structure=poscar.structure, iorder="c")
class _PotentialField(_Field):
"""Base class for potential-like fields."""
[docs]class VxcPotential(_PotentialField):
"""
XC Potential.
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: VxcPotential
"""
netcdf_name = "exchange_correlation_potential"
latex_label = "vxc $[eV/A^3]$"
[docs]class VhartreePotential(_PotentialField):
"""
Hartree Potential.
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: VhartreePotential
"""
netcdf_name = "vhartree"
latex_label = "vh $[eV/A^3]$"
[docs]class VhxcPotential(_PotentialField):
"""
Hartree + XC
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: VhxcPotential
"""
netcdf_name = "vhxc"
latex_label = "vhxc $[eV/A^3]$"
[docs]class VksPotential(_PotentialField):
"""
Hartree + XC + sum of local pseudo-potential terms.
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: VksPotential
"""
netcdf_name = "vtrial"
latex_label = "vks $[eV/A^3]$"
class FieldReader(ETSF_Reader):
"""
This object reads density data from a netcdf file.
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: FieldReader
"""
def read_den_dims(self):
"""Returns a |AttrDict| dictionary with the basic dimensions."""
return AttrDict(
nspinor=self.read_dimvalue("number_of_spinor_components"),
nsppol=self.read_dimvalue("number_of_spins"),
nspden=self.read_dimvalue("number_of_components"),
nfft1=self.read_dimvalue("number_of_grid_points_vector1"),
nfft2=self.read_dimvalue("number_of_grid_points_vector2"),
nfft3=self.read_dimvalue("number_of_grid_points_vector3"),
)
def read_field(self):
"""
Read and return the first field variable found in the netcdf_ file.
Raise:
`ValueError` if cannot find a field or multiple fields are found.
"""
found = [field_cls for field_cls in all_subclasses(_Field)
if field_cls.netcdf_name in self.rootgrp.variables]
if not found or len(found) > 1:
raise ValueError("Found `%s` fields in file: %s" % (str(found), self.path))
field_cls = found[0]
return self.read_denpot(varname=field_cls.netcdf_name, field_cls=field_cls)
def read_density(self, field_cls=Density):
"""Read density data. Return :class:`Density` object."""
return self.read_denpot(varname=field_cls.netcdf_name, field_cls=field_cls)
def read_vh(self, field_cls=VhartreePotential):
"""Read potential data. Return :class:`VhartreePotential` object."""
return self.read_denpot(varname=field_cls.netcdf_name, field_cls=field_cls)
def read_vxc(self, field_cls=VxcPotential):
"""Read potential data. Return :class:`VxcPotential` object."""
return self.read_denpot(varname=field_cls.netcdf_name, field_cls=field_cls)
def read_vhxc(self, field_cls=VhxcPotential):
"""Read potential data. Return :class:`VhxcPotential` object."""
return self.read_denpot(varname=field_cls.netcdf_name, field_cls=field_cls)
def read_vks(self, field_cls=VksPotential):
"""Read potential data. Return :class:`VksPotential` object."""
return self.read_denpot(varname=field_cls.netcdf_name, field_cls=field_cls)
#def read_vks1(self, field_cls=Vks1Potential):
# """Read potential data. Return :class:`Vks1Potential` object."""
# return self.read_denpot(varname=field_cls.netcdf_name, field_cls=field_cls)
def read_denpot(self, varname, field_cls):
"""
Factory function to read den/pot data from netcdf_ files and instantiate :class:`_Field` objects.
Note that unlike Abinit, datar[nspden] contains the up/down components if nsppol = 2
"""
structure = self.read_structure()
dims = self.read_den_dims()
# Abinit conventions:
# rhor(nfft, nspden) = electron density in real-space.
# (if spin polarized, array contains total density in first half and spin-up density in second half)
# (for non-collinear magnetism, first element: total density, 3 next ones: mx, my, mz in units of hbar/2)
datar = self.read_value(varname)
# Get rid of fake last dimensions if cplex == 1 or convert to complex array if cplex == 2.
# Remember that datar has been written from Fortran --> (n3, n2, n1).
cplex = datar.shape[-1]
if cplex == 1:
datar = np.reshape(datar, (dims.nspden, dims.nfft3, dims.nfft2, dims.nfft1))
else:
datar = datar[..., 0] + 1j * datar[..., 1]
if dims.nspinor == 1:
if dims.nspden == 1:
assert dims.nsppol == 1
elif dims.nspden == 2:
# assert dims.nsppol == 2
if issubclass(field_cls, _DensityField):
# If Density: store rho_up, rho_down instead of rho_total, rho_up.
total = datar[0].copy()
datar[0] = datar[1].copy()
datar[1] = total - datar[0]
else:
raise ValueError("Invalid nspinor: %s, nspden: %s and nsppol: %s" % (
dims.nspinor, dims.nspden, dims.nsppol))
#if issubclass(field_cls, _DensityField):
#elif issubclass(field_cls, _PotentialFieldField):
#else:
# raise TypeError("Don't know how to handle class: %s" % type(field_cls)
else:
if dims.nspden == 4:
raise NotImplementedError()
#datar_ab = np.empty_like(datar)
#datar_ab[0] = (datar[0] + datar[3]) / 2 # (n + mz) / 2
#datar_ab[1] = (rho[1] - 1j*rho[2]) / 2 # (mx - jmy) / 2
#datar_ab[2] = (rho[1] + 1j*rho[2]) / 2 # (mx + jmy) /2
#datar_ab[3] = (rho[0] - rho[3]) / 2 # (n - mz) / 2
#datar = datar_ab
else:
raise NotImplementedError()
# Structure uses Angstrom. Abinit uses Bohr.
if issubclass(field_cls, _DensityField):
fact = 1 / pmgu.bohr_to_angstrom ** 3
if issubclass(field_cls, _PotentialField):
fact = pmgu.Ha_to_eV / pmgu.bohr_to_angstrom ** 3
datar *= fact
# use iorder = "f" to transpose the last 3 dimensions since ETSF
# stores data in Fortran order while AbiPy uses C-ordering.
if cplex == 1:
return field_cls(dims.nspinor, dims.nsppol, dims.nspden, datar, structure, iorder="f")
else:
raise NotImplementedError("cplex %s not coded" % cplex)