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: Qiming Sun <osirpt.sun@gmail.com>
17#         Peter Koval <koval.peter@gmail.com>
18#         Paul J. Robinson <pjrobinson@ucla.edu>
19#
20
21'''
22Gaussian cube file format.  Reference:
23http://paulbourke.net/dataformats/cube/
24https://h5cube-spec.readthedocs.io/en/latest/cubeformat.html
25http://gaussian.com/cubegen/
26
27The output cube file has the following format
28
29Comment line
30Comment line
31N_atom Ox Oy Oz         # number of atoms, followed by the coordinates of the origin
32N1 vx1 vy1 vz1          # number of grids along each axis, followed by the step size in x/y/z direction.
33N2 vx2 vy2 vz2          # ...
34N3 vx3 vy3 vz3          # ...
35Atom1 Z1 x y z          # Atomic number, charge, and coordinates of the atom
36...                     # ...
37AtomN ZN x y z          # ...
38Data on grids           # (N1*N2) lines of records, each line has N3 elements
39'''
40
41import time
42import numpy
43import pyscf
44from pyscf import lib
45from pyscf import gto
46from pyscf import df
47from pyscf.dft import numint
48from pyscf import __config__
49
50RESOLUTION = getattr(__config__, 'cubegen_resolution', None)
51BOX_MARGIN = getattr(__config__, 'cubegen_box_margin', 3.0)
52ORIGIN = getattr(__config__, 'cubegen_box_origin', None)
53# If given, EXTENT should be a 3-element ndarray/list/tuple to represent the
54# extension in x, y, z
55EXTENT = getattr(__config__, 'cubegen_box_extent', None)
56
57
58def density(mol, outfile, dm, nx=80, ny=80, nz=80, resolution=RESOLUTION,
59            margin=BOX_MARGIN):
60    """Calculates electron density and write out in cube format.
61
62    Args:
63        mol : Mole
64            Molecule to calculate the electron density for.
65        outfile : str
66            Name of Cube file to be written.
67        dm : ndarray
68            Density matrix of molecule.
69
70    Kwargs:
71        nx : int
72            Number of grid point divisions in x direction.
73            Note this is function of the molecule's size; a larger molecule
74            will have a coarser representation than a smaller one for the
75            same value. Conflicts to keyword resolution.
76        ny : int
77            Number of grid point divisions in y direction.
78        nz : int
79            Number of grid point divisions in z direction.
80        resolution: float
81            Resolution of the mesh grid in the cube box. If resolution is
82            given in the input, the input nx/ny/nz have no effects.  The value
83            of nx/ny/nz will be determined by the resolution and the cube box
84            size.
85    """
86    from pyscf.pbc.gto import Cell
87    cc = Cube(mol, nx, ny, nz, resolution, margin)
88
89    GTOval = 'GTOval'
90    if isinstance(mol, Cell):
91        GTOval = 'PBC' + GTOval
92
93    # Compute density on the .cube grid
94    coords = cc.get_coords()
95    ngrids = cc.get_ngrids()
96    blksize = min(8000, ngrids)
97    rho = numpy.empty(ngrids)
98    for ip0, ip1 in lib.prange(0, ngrids, blksize):
99        ao = mol.eval_gto(GTOval, coords[ip0:ip1])
100        rho[ip0:ip1] = numint.eval_rho(mol, ao, dm)
101    rho = rho.reshape(cc.nx,cc.ny,cc.nz)
102
103    # Write out density to the .cube file
104    cc.write(rho, outfile, comment='Electron density in real space (e/Bohr^3)')
105    return rho
106
107
108def orbital(mol, outfile, coeff, nx=80, ny=80, nz=80, resolution=RESOLUTION,
109            margin=BOX_MARGIN):
110    """Calculate orbital value on real space grid and write out in cube format.
111
112    Args:
113        mol : Mole
114            Molecule to calculate the electron density for.
115        outfile : str
116            Name of Cube file to be written.
117        coeff : 1D array
118            coeff coefficient.
119
120    Kwargs:
121        nx : int
122            Number of grid point divisions in x direction.
123            Note this is function of the molecule's size; a larger molecule
124            will have a coarser representation than a smaller one for the
125            same value. Conflicts to keyword resolution.
126        ny : int
127            Number of grid point divisions in y direction.
128        nz : int
129            Number of grid point divisions in z direction.
130        resolution: float
131            Resolution of the mesh grid in the cube box. If resolution is
132            given in the input, the input nx/ny/nz have no effects.  The value
133            of nx/ny/nz will be determined by the resolution and the cube box
134            size.
135    """
136    from pyscf.pbc.gto import Cell
137    cc = Cube(mol, nx, ny, nz, resolution, margin)
138
139    GTOval = 'GTOval'
140    if isinstance(mol, Cell):
141        GTOval = 'PBC' + GTOval
142
143    # Compute density on the .cube grid
144    coords = cc.get_coords()
145    ngrids = cc.get_ngrids()
146    blksize = min(8000, ngrids)
147    orb_on_grid = numpy.empty(ngrids)
148    for ip0, ip1 in lib.prange(0, ngrids, blksize):
149        ao = mol.eval_gto(GTOval, coords[ip0:ip1])
150        orb_on_grid[ip0:ip1] = numpy.dot(ao, coeff)
151    orb_on_grid = orb_on_grid.reshape(cc.nx,cc.ny,cc.nz)
152
153    # Write out orbital to the .cube file
154    cc.write(orb_on_grid, outfile, comment='Orbital value in real space (1/Bohr^3)')
155    return orb_on_grid
156
157
158def mep(mol, outfile, dm, nx=80, ny=80, nz=80, resolution=RESOLUTION,
159        margin=BOX_MARGIN):
160    """Calculates the molecular electrostatic potential (MEP) and write out in
161    cube format.
162
163    Args:
164        mol : Mole
165            Molecule to calculate the electron density for.
166        outfile : str
167            Name of Cube file to be written.
168        dm : ndarray
169            Density matrix of molecule.
170
171    Kwargs:
172        nx : int
173            Number of grid point divisions in x direction.
174            Note this is function of the molecule's size; a larger molecule
175            will have a coarser representation than a smaller one for the
176            same value. Conflicts to keyword resolution.
177        ny : int
178            Number of grid point divisions in y direction.
179        nz : int
180            Number of grid point divisions in z direction.
181        resolution: float
182            Resolution of the mesh grid in the cube box. If resolution is
183            given in the input, the input nx/ny/nz have no effects.  The value
184            of nx/ny/nz will be determined by the resolution and the cube box
185            size.
186    """
187    cc = Cube(mol, nx, ny, nz, resolution, margin)
188
189    coords = cc.get_coords()
190
191    # Nuclear potential at given points
192    Vnuc = 0
193    for i in range(mol.natm):
194        r = mol.atom_coord(i)
195        Z = mol.atom_charge(i)
196        rp = r - coords
197        Vnuc += Z / numpy.einsum('xi,xi->x', rp, rp)**.5
198
199    # Potential of electron density
200    Vele = numpy.empty_like(Vnuc)
201    for p0, p1 in lib.prange(0, Vele.size, 600):
202        fakemol = gto.fakemol_for_charges(coords[p0:p1])
203        ints = df.incore.aux_e2(mol, fakemol)
204        Vele[p0:p1] = numpy.einsum('ijp,ij->p', ints, dm)
205
206    MEP = Vnuc - Vele     # MEP at each point
207    MEP = MEP.reshape(cc.nx,cc.ny,cc.nz)
208
209    # Write the potential
210    cc.write(MEP, outfile, 'Molecular electrostatic potential in real space')
211    return MEP
212
213
214class Cube(object):
215    '''  Read-write of the Gaussian CUBE files
216
217    Attributes:
218        nx : int
219            Number of grid point divisions in x direction.
220            Note this is function of the molecule's size; a larger molecule
221            will have a coarser representation than a smaller one for the
222            same value. Conflicts to keyword resolution.
223        ny : int
224            Number of grid point divisions in y direction.
225        nz : int
226            Number of grid point divisions in z direction.
227        resolution: float
228            Resolution of the mesh grid in the cube box. If resolution is
229            given in the input, the input nx/ny/nz have no effects.  The value
230            of nx/ny/nz will be determined by the resolution and the cube box
231            size. The unit is Bohr.
232    '''
233    def __init__(self, mol, nx=80, ny=80, nz=80, resolution=RESOLUTION,
234                 margin=BOX_MARGIN, origin=ORIGIN, extent=EXTENT):
235        from pyscf.pbc.gto import Cell
236        self.mol = mol
237        coord = mol.atom_coords()
238
239        # If the molecule is periodic, use lattice vectors as the box
240        # and automatically determine margin, origin, and extent
241        if isinstance(mol, Cell):
242            self.box = mol.lattice_vectors()
243            atom_center = (numpy.max(coord, axis=0) + numpy.min(coord, axis=0))/2
244            box_center = (self.box[0] + self.box[1] + self.box[2])/2
245            self.boxorig = atom_center - box_center
246        else:
247            # Create a box based on its origin and extents/lengths (rectangular cuboid only)
248            # If extent is not supplied, use the coordinates plus a margin on both sides
249            if extent is None:
250                extent = numpy.max(coord, axis=0) - numpy.min(coord, axis=0) + 2*margin
251            self.box = numpy.diag(extent)
252
253            # if origin is not supplied, set it as the minimum coordinate minus a margin.
254            if origin is None:
255                origin = numpy.min(coord, axis=0) - margin
256            self.boxorig = numpy.asarray(origin)
257
258        if resolution is not None:
259            nx, ny, nz = numpy.ceil(numpy.diag(self.box) / resolution).astype(int)
260
261        self.nx = nx
262        self.ny = ny
263        self.nz = nz
264
265        if isinstance(mol, Cell):
266            # Use an asymmetric mesh for tiling unit cells
267            self.xs = numpy.linspace(0, 1, nx, endpoint=False)
268            self.ys = numpy.linspace(0, 1, ny, endpoint=False)
269            self.zs = numpy.linspace(0, 1, nz, endpoint=False)
270        else:
271            # Use endpoint=True to get a symmetric mesh
272            # see also the discussion https://github.com/sunqm/pyscf/issues/154
273            self.xs = numpy.linspace(0, 1, nx, endpoint=True)
274            self.ys = numpy.linspace(0, 1, ny, endpoint=True)
275            self.zs = numpy.linspace(0, 1, nz, endpoint=True)
276
277    def get_coords(self) :
278        """  Result: set of coordinates to compute a field which is to be stored
279        in the file.
280        """
281        frac_coords = lib.cartesian_prod([self.xs, self.ys, self.zs])
282        return frac_coords @ self.box + self.boxorig # Convert fractional coordinates to real-space coordinates
283
284    def get_ngrids(self):
285        return self.nx * self.ny * self.nz
286
287    def get_volume_element(self):
288        return (self.xs[1]-self.xs[0])*(self.ys[1]-self.ys[0])*(self.zs[1]-self.zs[0])
289
290    def write(self, field, fname, comment=None):
291        """  Result: .cube file with the field in the file fname.  """
292        assert(field.ndim == 3)
293        assert(field.shape == (self.nx, self.ny, self.nz))
294        if comment is None:
295            comment = 'Generic field? Supply the optional argument "comment" to define this line'
296
297        mol = self.mol
298        coord = mol.atom_coords()
299        with open(fname, 'w') as f:
300            f.write(comment+'\n')
301            f.write(f'PySCF Version: {pyscf.__version__}  Date: {time.ctime()}\n')
302            f.write(f'{mol.natm:5d}')
303            f.write('%12.6f%12.6f%12.6f\n' % tuple(self.boxorig.tolist()))
304            dx = self.xs[-1] if len(self.xs) == 1 else self.xs[1]
305            dy = self.ys[-1] if len(self.ys) == 1 else self.ys[1]
306            dz = self.zs[-1] if len(self.zs) == 1 else self.zs[1]
307            delta = (self.box.T * [dx,dy,dz]).T
308            f.write(f'{self.nx:5d}{delta[0,0]:12.6f}{delta[0,1]:12.6f}{delta[0,2]:12.6f}\n')
309            f.write(f'{self.ny:5d}{delta[1,0]:12.6f}{delta[1,1]:12.6f}{delta[1,2]:12.6f}\n')
310            f.write(f'{self.nz:5d}{delta[2,0]:12.6f}{delta[2,1]:12.6f}{delta[2,2]:12.6f}\n')
311            for ia in range(mol.natm):
312                atmsymb = mol.atom_symbol(ia)
313                f.write('%5d%12.6f'% (gto.charge(atmsymb), 0.))
314                f.write('%12.6f%12.6f%12.6f\n' % tuple(coord[ia]))
315
316            for ix in range(self.nx):
317                for iy in range(self.ny):
318                    for iz0, iz1 in lib.prange(0, self.nz, 6):
319                        fmt = '%13.5E' * (iz1-iz0) + '\n'
320                        f.write(fmt % tuple(field[ix,iy,iz0:iz1].tolist()))
321
322    def read(self, cube_file):
323        with open(cube_file, 'r') as f:
324            f.readline()
325            f.readline()
326            data = f.readline().split()
327            natm = int(data[0])
328            self.boxorig = numpy.array([float(x) for x in data[1:]])
329            def parse_nx(data):
330                d = data.split()
331                return int(d[0]), numpy.array([float(x) for x in d[1:]])
332            self.nx, self.xs = parse_nx(f.readline())
333            self.ny, self.ys = parse_nx(f.readline())
334            self.nz, self.zs = parse_nx(f.readline())
335            atoms = []
336            for ia in range(natm):
337                d = f.readline().split()
338                atoms.append([int(d[0]), [float(x) for x in d[2:]]])
339            self.mol = gto.M(atom=atoms, unit='Bohr')
340
341            data = f.read()
342        cube_data = numpy.array([float(x) for x in data.split()])
343        return cube_data.reshape([self.nx, self.ny, self.nz])
344
345
346if __name__ == '__main__':
347    from pyscf import scf
348    from pyscf.tools import cubegen
349    mol = gto.M(atom='''O 0.00000000,  0.000000,  0.000000
350                H 0.761561, 0.478993, 0.00000000
351                H -0.761561, 0.478993, 0.00000000''', basis='6-31g*')
352    mf = scf.RHF(mol).run()
353    cubegen.density(mol, 'h2o_den.cube', mf.make_rdm1()) #makes total density
354    cubegen.mep(mol, 'h2o_pot.cube', mf.make_rdm1())
355    cubegen.orbital(mol, 'h2o_mo1.cube', mf.mo_coeff[:,0])
356