fields.py

# coding: utf-8
"""This module contains the class describing densities in real space on uniform 3D meshes."""
from __future__ import print_function, division, unicode_literals, absolute_import

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
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 transpose_last3dims, duck
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,x,y) --> (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:
                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(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[1]))

        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
            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))

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$]"

    @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.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:
                    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:
                            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')

    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])

    @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.")

    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")

    @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)

    @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:]

    @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)

    @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

    @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

    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)

    @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

    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

    @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."""


class VxcPotential(_PotentialField):
    """
    XC Potential.

    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: VxcPotential
    """
    netcdf_name = "exchange_correlation_potential"
    latex_label = "vxc $[eV/A^3]$"


class VhartreePotential(_PotentialField):
    """
    Hartree Potential.

    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: VhartreePotential
    """
    netcdf_name = "vhartree"
    latex_label = "vh $[eV/A^3]$"


class VhxcPotential(_PotentialField):
    """
    Hartree + XC

    .. rubric:: Inheritance Diagram
    .. inheritance-diagram:: VhxcPotential
    """
    netcdf_name = "vhxc"
    latex_label = "vhxc $[eV/A^3]$"


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_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)