tryout_pyscf.py

"""
Simple script to learn the basics of PySCF. This script computes
the optimal geometry of a molecule, and then visualizes its electron
density.

Commands for installing the necessary modules:

# Install PySCF
python3 -m pip install pyscf

# Prerequisite modules in order to get geometry optimization to work.
python3 -m pip install pyberny geometric
python3 -m pip install git+https://github.com/pyscf/qsdopt

# Use Plotly for 3D visualizations.
python3 -m pip install plotly

Reference:

    Recent developments in the PySCF program package,
    Qiming Sun, Xing Zhang, Samragni Banerjee, Peng Bao,
    Marc Barbry, Nick S. Blunt, Nikolay A. Bogdanov,
    George H. Booth, Jia Chen, Zhi-Hao Cui, Janus J. Eriksen,
    Yang Gao, Sheng Guo, Jan Hermann, Matthew R. Hermes,
    Kevin Koh, Peter Koval, Susi Lehtola, Zhendong Li,
    Junzi Liu, Narbe Mardirossian, James D. McClain,
    Mario Motta, Bastien Mussard, Hung Q. Pham, Artem Pulkin,
    Wirawan Purwanto, Paul J. Robinson, Enrico Ronca, 
    Elvira R. Sayfutyarova, Maximilian Scheurer, Henry F. Schurkus,
    James E. T. Smith, Chong Sun, Shi-Ning Sun, Shiv Upadhyay,
    Lucas K. Wagner, Xiao Wang, Alec White, James Daniel Whitfield,
    Mark J. Williamson, Sebastian Wouters, Jun Yang, Jason M. Yu,
    Tianyu Zhu, Timothy C. Berkelbach, Sandeep Sharma,
    Alexander Yu. Sokolov, and Garnet Kin-Lic Chan,
    J. Chem. Phys., 153, 024109 (2020). doi:10.1063/5.0006074

"""
from pyscf import gto, dft
# from pyscf import scf, cc
# from pyscf.tools import chgcar
import numpy as np
import plotly.graph_objects as go
from pyscf.geomopt import geometric_solver


molecule = gto.Mole()
molecule.build(
    atom='H 0 0 1.5; O 0 0 0; H 0 1.5 0',
    basis='aug-cc-pVTZ',
    unit='bohr'
)
print(molecule.atom_coords())
rks = dft.RKS(molecule)
# rks.check_convergence()
# rks.conv_tol = 1e-16
molecule = geometric_solver.optimize(rks, max_steps=200)
# res = cc.CCSD(res).run()

res = rks.run()
print(molecule.atom_coords())

N = 128
X, Y, Z = np.meshgrid(*3*[np.linspace(-3.0, 3.0, 128)])
coords = np.array([X, Y, Z]).reshape(3, N*N*N).T
density_matrix = res.make_rdm1()
ao = dft.numint.eval_ao(molecule, coords)
electron_density = dft.numint.eval_rho(
    molecule, ao, density_matrix).reshape(N, N, N)
print(electron_density.shape)

fig = go.Figure(data=go.Volume(
    x=X.flatten(), y=Y.flatten(), z=Z.flatten(),
    value=electron_density.flatten(), isomin=0.1, isomax=2.0,
    opacity=0.1, surface_count=20,  
))
fig.show()