1#!/usr/bin/env python
2# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16# Author: Paul J. Robinson <pjrobinson@ucla.edu>
17#         Qiming Sun <osirpt.sun@gmail.com>
18#
19
20'''
21Vasp CHGCAR file format
22
23See also
24https://cms.mpi.univie.ac.at/vasp/vasp/CHGCAR_file.html
25'''
26
27import sys
28import collections
29import time
30import numpy
31import pyscf
32from pyscf import lib
33from pyscf import gto
34from pyscf.pbc import gto as pbcgto
35from pyscf.tools import cubegen
36
37if sys.version_info >= (3,):
38    unicode = str
39
40RESOLUTION = cubegen.RESOLUTION
41BOX_MARGIN = cubegen.BOX_MARGIN
42
43
44def density(cell, outfile, dm, nx=60, ny=60, nz=60, resolution=RESOLUTION):
45    '''Calculates electron density and write out in CHGCAR format.
46
47    Args:
48        cell : Mole or Cell object
49            Mole or pbc Cell. If Mole object is given, the program will guess
50            a cubic lattice for the molecule.
51        outfile : str
52            Name of Cube file to be written.
53        dm : ndarray
54            Density matrix of molecule.
55
56    Kwargs:
57        nx : int
58            Number of grid point divisions in x direction.
59            Note this is function of the molecule's size; a larger molecule
60            will have a coarser representation than a smaller one for the
61            same value.
62        ny : int
63            Number of grid point divisions in y direction.
64        nz : int
65            Number of grid point divisions in z direction.
66
67    Returns:
68        No return value. This function outputs a VASP chgcarlike file
69        (with phase if desired)...it can be opened in VESTA or VMD or
70        many other softwares
71
72    Examples:
73
74        >>> # generates the first MO from the list of mo_coefficents
75        >>> from pyscf.pbc import gto, scf
76        >>> from pyscf.tools import chgcar
77        >>> cell = gto.M(atom='H 0 0 0; H 0 0 1', a=numpy.eye(3)*3)
78        >>> mf = scf.RHF(cell).run()
79        >>> chgcar.density(cell, 'h2.CHGCAR', mf.make_rdm1())
80
81    '''
82    cc = CHGCAR(cell, nx=nx, ny=ny, nz=nz, resolution=resolution)
83
84    coords = cc.get_coords()
85    ngrids = cc.get_ngrids()
86    blksize = min(8000, ngrids)
87    rho = numpy.empty(ngrids)
88    for ip0, ip1 in lib.prange(0, ngrids, blksize):
89        if isinstance(cell, pbcgto.cell.Cell):
90            ao = cell.pbc_eval_gto('GTOval', coords[ip0:ip1])
91        else:
92            ao = cell.eval_gto('GTOval', coords[ip0:ip1])
93        rho[ip0:ip1] = lib.einsum('pi,ij,pj->p', ao, dm, ao)
94    rho = rho.reshape(nx,ny,nz)
95
96    cc.write(rho, outfile)
97
98
99def orbital(cell, outfile, coeff, nx=60, ny=60, nz=60, resolution=RESOLUTION):
100    '''Calculate orbital value on real space grid and write out in
101    CHGCAR format.
102
103    Args:
104        cell : Mole or Cell object
105            Mole or pbc Cell. If Mole object is given, the program will guess
106            a cubic lattice for the molecule.
107        outfile : str
108            Name of Cube file to be written.
109        dm : ndarray
110            Density matrix of molecule.
111
112    Kwargs:
113        nx : int
114            Number of grid point divisions in x direction.
115            Note this is function of the molecule's size; a larger molecule
116            will have a coarser representation than a smaller one for the
117            same value.
118        ny : int
119            Number of grid point divisions in y direction.
120        nz : int
121            Number of grid point divisions in z direction.
122
123    Returns:
124        No return value. This function outputs a VASP chgcarlike file
125        (with phase if desired)...it can be opened in VESTA or VMD or
126        many other softwares
127
128    Examples:
129
130        >>> # generates the first MO from the list of mo_coefficents
131        >>> from pyscf.pbc import gto, scf
132        >>> from pyscf.tools import chgcar
133        >>> cell = gto.M(atom='H 0 0 0; H 0 0 1', a=numpy.eye(3)*3)
134        >>> mf = scf.RHF(cell).run()
135        >>> chgcar.orbital(cell, 'h2_mo1.CHGCAR', mf.mo_coeff[:,0])
136
137    '''
138    cc = CHGCAR(cell, nx=nx, ny=ny, nz=nz, resolution=resolution)
139
140    coords = cc.get_coords()
141    ngrids = cc.get_ngrids()
142    blksize = min(8000, ngrids)
143    orb_on_grid = numpy.empty(ngrids)
144    for ip0, ip1 in lib.prange(0, ngrids, blksize):
145        if isinstance(cell, pbcgto.cell.Cell):
146            ao = cell.pbc_eval_gto('GTOval', coords[ip0:ip1])
147        else:
148            ao = cell.eval_gto('GTOval', coords[ip0:ip1])
149        orb_on_grid[ip0:ip1] = numpy.dot(ao, coeff)
150    orb_on_grid = orb_on_grid.reshape(nx,ny,nz)
151
152    cc.write(orb_on_grid, outfile, comment='Orbital value in real space (1/Bohr^3)')
153
154
155class CHGCAR(cubegen.Cube):
156    '''  Read-write of the Vasp CHGCAR files  '''
157    def __init__(self, cell, nx=60, ny=60, nz=60, resolution=RESOLUTION,
158                 margin=BOX_MARGIN):
159        if not isinstance(cell, pbcgto.cell.Cell):
160            coord = cell.atom_coords()
161            box = numpy.max(coord,axis=0) - numpy.min(coord,axis=0) + margin*2
162            boxorig = numpy.min(coord,axis=0) - margin
163            if resolution is not None:
164                nx, ny, nz = numpy.ceil(box / resolution).astype(int)
165            self.box = numpy.diag(box)
166            lib.logger.warn(cell, 'Molecular system is found. FFT-grid is not '
167                            'available for Molecule. Lattice (in Bohr)\n'
168                            '%s\nand FFT grids %s are applied.',
169                            self.box, (nx,ny,nz))
170            self.mol = cell
171            cell = cell.view(pbcgto.Cell)
172            if (isinstance(cell.unit, (str, unicode)) and
173                cell.unit.startswith(('B','b','au','AU'))):
174                cell.a = self.box
175            else:
176                cell.a = self.box * lib.param.BOHR
177            ptr = cell._atm[:,gto.PTR_COORD]
178            cell._env[ptr+0] = coord[:,0] - boxorig[0]
179            cell._env[ptr+1] = coord[:,1] - boxorig[1]
180            cell._env[ptr+2] = coord[:,2] - boxorig[2]
181
182        self.nx = nx
183        self.ny = ny
184        self.nz = nz
185        self.cell = cell
186        self.box = cell.lattice_vectors()
187        self.boxorig = numpy.zeros(3)
188        self.vol = cell.vol
189
190    def get_coords(self) :
191        """  Result: set of coordinates to compute a field which is to be stored
192        in the file.
193        """
194        xs = numpy.arange(self.nx) * (1./self.nx)
195        ys = numpy.arange(self.ny) * (1./self.ny)
196        zs = numpy.arange(self.nz) * (1./self.nz)
197        xyz = lib.cartesian_prod((xs, ys, zs))
198        coords = numpy.dot(xyz, self.box)
199        return numpy.asarray(coords, order='C')
200
201    def write(self, field, fname, comment=None):
202        """  Result: .vasp file with the field in the file fname.  """
203        assert(field.ndim == 3)
204        assert(field.shape == (self.nx, self.ny, self.nz))
205        if comment is None:
206            comment = 'VASP file: Electron density in real space (e/Bohr^3)  '
207
208        cell = self.cell
209
210        # See CHGCAR format https://cms.mpi.univie.ac.at/vasp/vasp/CHGCAR_file.html
211        # the value of (total density * volume) was dumped
212        field = field * self.vol
213
214        boxA = self.box * lib.param.BOHR
215        atomList= [cell.atom_pure_symbol(i) for i in range(cell.natm)]
216        Axyz = zip(atomList, cell.atom_coords().tolist())
217        Axyz = sorted(Axyz, key = lambda x: x[0])
218        swappedCoords = [(vec[1]+self.boxorig) * lib.param.BOHR for vec in Axyz]
219        vaspAtomicInfo = collections.Counter([xyz[0] for xyz in Axyz])
220        vaspAtomicInfo = sorted(vaspAtomicInfo.items())
221        with open(fname, 'w') as f:
222            f.write(comment)
223            f.write('PySCF Version: %s  Date: %s\n' % (pyscf.__version__, time.ctime()))
224            f.write('1.0000000000\n')
225            f.write('%14.8f %14.8f %14.8f \n' % (boxA[0,0],boxA[0,1],boxA[0,2]))
226            f.write('%14.8f %14.8f %14.8f \n' % (boxA[1,0],boxA[1,1],boxA[1,2]))
227            f.write('%14.8f %14.8f %14.8f \n' % (boxA[2,0],boxA[2,1],boxA[2,2]))
228            f.write(''.join(['%5.3s'%atomN[0] for atomN in vaspAtomicInfo]) + '\n')
229            f.write(''.join(['%5d'%atomN[1] for atomN in vaspAtomicInfo]) + '\n')
230            f.write('Cartesian \n')
231            for ia in range(cell.natm):
232                f.write(' %14.8f %14.8f %14.8f\n' % tuple(swappedCoords[ia]))
233            f.write('\n')
234            f.write('%6.5s %6.5s %6.5s \n' % (self.nx,self.ny,self.nz))
235            fmt = ' %14.8e '
236            for iz in range(self.nz):
237                for iy in range(self.ny):
238                    f.write('\n')
239                    for ix in range(self.nx):
240                        f.write(fmt % field[ix,iy,iz])
241
242    def read(self, chgcar_file):
243        raise NotImplementedError
244
245
246if __name__ == '__main__':
247    from pyscf.pbc import scf
248    from pyscf.tools import chgcar
249    cell = gto.M(atom='H 0 0 0; H 0 0 1', a=numpy.eye(3)*3)
250    mf = scf.RHF(cell).run()
251    chgcar.density(cell, 'h2.CHGCAR', mf.make_rdm1()) #makes total density
252    chgcar.orbital(cell, 'h2_mo1.CHGCAR', mf.mo_coeff[:,0]) # makes mo#1 (sigma)
253    chgcar.orbital(cell, 'h2_mo2.CHGCAR', mf.mo_coeff[:,1]) # makes mo#2 (sigma*)
254
255