Source code for nexus.pyscf_input
##################################################################
## (c) Copyright 2018- by Jaron T. Krogel ##
##################################################################
#====================================================================#
# quantum_package_input.py #
# Supports read/write of PySCF input. #
# #
# Content summary: #
# PyscfInput #
# SimulationInput class for PySCF. #
# #
# generate_pyscf_input #
# User-facing function to create template based input. #
#====================================================================#
import numpy as np
from .developer import obj, error
from .simulation import SimulationInputTemplateDev
[docs]
def render_string(s_in,n):
indent = n*' '
s = ''
lines = s_in.splitlines()
s += 3*"'"+'\n'
for line in lines:
s += indent+line.strip()+'\n'
#end for
s += indent+3*"'"
return s
#end def render_string
[docs]
def render_array(a,n):
indent = n*' '
s = ''
if len(a.shape)==1:
s += 'array('+repr(list(a))+')'
elif len(a.shape)==2:
s += 'array([\n'
for a1 in a:
s += indent+repr(list(a1))+',\n'
#end for
s = s[:-2]+'])'
else:
error('render_array only supports up to 2D arrays')
#end if
return s
#end def render_array
[docs]
class PyscfInput(SimulationInputTemplateDev):
basic_types = bool,int,float,str,tuple,list,dict
allowed_types = basic_types+(np.ndarray,)
mole_order = '''
dump_input
parse_arg
verbose
output
max_memory
atom
basis
unit
nucmod
ecp
charge
spin
symmetry
symmetry_subgroup
cart
nelec
nelectron
multiplicity
ms
'''.split()
cell_order = '''
dump_input
parse_arg
a
mesh
ke_cutoff
precision
nimgs
ew_eta
ew_cut
pseudo
basis
h
dimension
rcut
ecp
low_dim_ft_type
unit
atom
gs
h
drop_exponent
nimgs
'''.split()
for k in mole_order:
if k not in cell_order:
cell_order.append(k)
#end if
#end for
mole_allowed = set(mole_order)
cell_allowed = set(cell_order)
def __init__(self,
template = None, # path to template input file
prefix = None, # $prefix var for file prefixes
custom = None, # obj w/ $ prefixed vars in template
system = None, # physical system object
units = None, # input units desired
use_folded = True, # use folded system/primitive cell
mole = None, # obj w/ Mole variables
cell = None, # obj w/ Cell variables
sys_var = None, # local var name for Mole/Cell
mole_var = 'mol', # local var name for Mole in written input
cell_var = 'cell', # local var name for Cell in written input
save_qmc = False, # convert to QMCPACK format
checkpoint = False, # set $chkfile variable
mf_var = 'mf', # local var name for mf, used for convert
kpts_var = 'kpts', # local var name for kpts, used for convert
filepath = None, # alias for template
text = None, # full text of (and alternate to) template
calculation = None, # obj w/ Calculation variables
chkfile = None, # obj w/ Calculation variables
twist_num = None, # Twist index
python_exe = 'python3', # Python executable
):
if filepath is None and template is not None:
filepath = template
elif calculation is not None:
self.calculation = calculation
text='''#! /usr/bin/env $python_exe
import numpy as np
from numpy import array
$pyscfimport
$system
$calculation
'''
#end if
SimulationInputTemplateDev.__init__(self,filepath,text)
self.prefix = prefix
self.save_qmc = save_qmc
self.checkpoint = checkpoint
self.addendum = None # used for save2qmcpack
if custom is not None:
self.assign(**custom)
#end if
if sys_var is not None:
mole_var = sys_var
cell_var = sys_var
#end if
sys_name = None
sys_inputs = obj()
sys_allowed = None
sys_kpoints = None
is_mole = mole is not None
is_cell = cell is not None
if is_mole and is_cell:
self.error('both mole and cell provided\nplease provide only one of them\nsystem cannot be both molecule and periodic solid')
elif is_mole and not isinstance(mole,(dict,obj)):
self.error('mole input must be a dict or obj\nyou provided input of type: {0}'.format(mole.__class__.__name__))
elif is_cell and not isinstance(cell,(dict,obj)):
self.error('cell input must be a dict or obj\nyou provided input of type: {0}'.format(cell.__class__.__name__))
#end if
extra = ''
if filepath is not None:
extra = '\ntemplate located at: {0}'.format(filepath)
#end if
if system is not None and 'system' not in self.keywords:
self.error('system input is provided, but $system is not present in template input'+extra)
#end if
folded_structure = None
tiled_structure = None
if system is not None and 'system' not in self.values:
if 'system' not in self.keywords:
self.error('cannot incorporate "system" input\n$system is not present in template input'+extra)
#end if
system = system.copy() # make a local copy
if use_folded and system.has_folded():
tiled_structure = system.structure.copy()
system = system.folded_system
#end if
s = system.structure
if units is not None:
s.change_units(units)
#end if
is_solid = s.has_axes()
is_molecule = not is_solid
if is_solid and is_mole:
self.error('mole input provided, but provided system is not a molecule (cell axes are present)')
elif is_molecule and is_cell:
self.error('cell input provided, but provided system is a molecule (cell axes are not present)')
#end if
is_mole |= is_molecule
is_cell |= is_solid
sys_inputs.atom = s.pos_to_str(units=s.units, with_elem=True)
sys_inputs.unit = s.units
sys_inputs.charge = system.net_charge
sys_inputs.spin = system.net_spin
if is_solid:
folded_structure = s.copy()
sys_inputs.dimension = len(s.axes)
sys_inputs.a = s.write_axes()
if len(s.kpoints)>0:
skp = s.copy()
skp.change_units('B')
sys_kpoints = skp.kpoints.copy()
#end if
#end if
#end if
if is_mole:
sys_name = 'mole'
sys_var = mole_var
sys_allowed = PyscfInput.mole_allowed
sys_order = PyscfInput.mole_order
if mole is not None:
sys_inputs.set(**mole)
#end if
elif is_cell:
sys_name = 'cell'
sys_var = cell_var
sys_allowed = PyscfInput.cell_allowed
sys_order = PyscfInput.cell_order
if cell is not None:
sys_inputs.set(**cell)
#end if
else:
None # no action needed if not molecule or periodic solid
#end if
if calculation is not None and 'calculation' not in self.values:
calc = calculation.copy() # make a local copy
calc.set_optional(
method = 'RKS',
df_fitting = True,
xc = 'pbe',
tol = '1e-10',
df_method = 'GDF',
exxdiv = 'ewald',
u_idx = None,
max_cycle = None,
level_shift = None,
chkfile = None,
u_val = None,
C_ao_lo = 'minao',
)
if calc.u_val is not None:
calc.u_val = np.array(calc.u_val)
#end if
if calc.u_idx is not None:
calc.u_idx = np.array(calc.u_idx)
#end if
c = '\n### generated calculation text ###\n'
if sys_name is not None:
df_str = '.density_fit()' if calc.df_fitting else ''
if sys_name == 'mole':
if calc.u_idx is None:
c += 'mf = scf.{}({}){}\n'.format(calc.method,sys_var,df_str)
else:
c += 'mf = dft.{}({},U_idx={},U_val={},C_ao_lo=\'{}\'){}\n'.format(calc.method,sys_var,render_array(calc.u_idx,1),render_array(calc.u_val,1),calc.C_ao_lo,df_str)
#end if
elif sys_name == 'cell':
c += 'mydf = df.{}({})\n'.format(calc.df_method,sys_var,'kpts')
c += 'mydf.auxbasis = \'weigend\'\n'
c += 'dfpath = \'df_ints.h5\'\n'
c += 'mydf._cderi_to_save = dfpath\n'
c += 'mydf.build()\n\n'
if calc.u_idx is None:
c += 'mf = scf.{}({},{}){}\n'.format(calc.method,sys_var,'kpts',df_str)
else:
c += 'mf = dft.{}({},{},U_idx={},U_val={},C_ao_lo=\'{}\'){}\n'.format(calc.method,sys_var,'kpts',render_array(calc.u_idx,1),render_array(calc.u_val,1),calc.C_ao_lo,df_str)
c += 'mf.exxdiv = \'{}\'\n'.format(calc.exxdiv)
#end if
if calc.max_cycle is not None:
c += 'mf.max_cycle={}\n'.format(calc.max_cycle)
#end if
if calc.level_shift is not None:
c += 'mf.level_shift={}\n'.format(calc.level_shift)
#end if
if calc.chkfile is not None:
c += 'mf.chkfile=\'{}\'\n'.format(calc.chkfile)
#end if
#end if
if 'KS' in calc.method:
c += 'mf.xc = \'{}\'\n'.format(calc.xc)
#end if
c += 'mf.tol = \'{}\'\n'.format(calc.tol)
if calc.df_fitting and not is_mole:
c += 'mf.with_df = mydf\n'
#end if
c += 'e_scf = mf.kernel()\n'
c += '### end generated calculation text ###\n\n'
self.assign(calculation=c)
cimp = '\n### generated pyscfimport text ###\n'
if is_mole:
cimp += 'from pyscf import df, scf, dft\n'
elif is_cell:
cimp += 'from pyscf.pbc import df, scf\n'
#endif
cimp += '### end generated pyscfimport text ###\n\n'
self.assign(pyscfimport=cimp)
self.assign(python_exe=python_exe)
#end if
if sys_name is not None:
invalid = set(sys_inputs.keys())-sys_allowed
if len(invalid)>0:
self.error('invalid {0} inputs\ninvalid inputs: {1}\nvalid options are: {2}'.format(sys_name,sorted(invalid),sorted(sys_allowed)))
#end if
klen = 0
has_array = False
for k,v in sys_inputs.items():
if not isinstance(v,PyscfInput.allowed_types):
tlist = ''
for t in PyscfInput.allowed_types:
tlist += t.__name__+','
#end for
tlist = tlist[:-1]
self.error('{0} input "{1}" has an invalid type\ninvalid type: {2}\nallowed types are: {3}'.format(sys_name,k,v.__class__.__name__,tlist))
#end if
klen = max(klen,len(k))
has_array |= isinstance(v,np.ndarray)
#end for
has_array |= isinstance(sys_kpoints,np.ndarray)
c = '\n### generated system text ###\n'
if has_array:
c += 'from numpy import array\n'
#end if
if is_mole:
c += 'from pyscf import gto as gto_loc\n'
c += '{0} = gto_loc.Mole()\n'.format(sys_var)
elif is_cell:
c += 'from pyscf.pbc import gto as gto_loc\n'
c += '{0} = gto_loc.Cell()\n'.format(sys_var)
#end if
fmt = sys_var+'.{0:<'+str(klen)+'} = {1}\n'
nalign = klen+len(sys_var)+4
for k in sys_order:
if k in sys_inputs:
v = sys_inputs[k]
if isinstance(v,str) and '\n' in v:
vs = render_string(v,nalign)
elif isinstance(v,np.ndarray):
if len(v.shape)>2:
self.error('cannot write system input variable {0}\n{0} is an array with more than two dimensions\nonly two dimensions are currently supported for writing\narray contents: {1}'.format(k,v))
#end if
vs = render_array(v,nalign)
elif isinstance(v,PyscfInput.basic_types):
vs = repr(v)
else:
vs = None # should not be possible
#end if
c += fmt.format(k,vs)
#end if
#end for
c += '{0}.build()\n'.format(sys_var)
if sys_kpoints is not None:
c += '{0} = {1}\n'.format(kpts_var,render_array(sys_kpoints,4))
#end if
c += '### end generated system text ###\n\n'
self.assign(system=c)
#end if
if prefix is not None:
self.allow_no_assign('prefix')
self.assign(prefix=prefix)
#end if
if save_qmc:
if prefix is None:
self.error('cannot generate save2qmcpack text\nplease provide input variable "prefix"\n(used to set "title" in save2qmcpack)')
elif sys_var is None:
self.error('cannot generate save2qmcpack text\nplease provide input variable "sys_var"\n(used to set "cell" in save2qmcpack) ')
#end if
if folded_structure is not None:
folded_structure.change_units('B') # always use Bohr units for k-points
nkpoints = len(folded_structure.kpoints)
if sys_kpoints is None and nkpoints!=0 or sys_kpoints is not None and len(sys_kpoints)!=nkpoints:
self.error('inconsistency in written and saved k-points')
#end if
if tiled_structure is None:
tiling = np.array([1,1,1],dtype=int)
ntwists = nkpoints
else:
if nkpoints==0:
self.error('k-points must be present for save2qmcpack to write supercell wavefunction')
#end if
tiled_structure.change_units('B') # always use Bohr units for twists/k-points
tiled_kpoints = tiled_structure.kpoints.copy()
ntwists = len(tiled_kpoints)
if ntwists==0:
self.error('supercell k-points must be present for save2qmcpack to write supercell wavefunction')
#end if
kmap = tiled_structure.kmap()
if len(kmap)!=ntwists:
self.error('inconsistency between supercell twist mapping and supercell twist count')
#end if
assert len(kmap)==ntwists
kmap_array = []
folded_indices = []
for itwist in sorted(kmap.keys()):
ifolded = list(kmap[itwist])
kmap_array.append(ifolded)
folded_indices.extend(ifolded)
#end for
kmap_array = np.array(kmap_array,dtype=int)
if len(set(folded_indices))!=nkpoints:
self.error('inconsistency in mapping between supercell twists and folded cell k-points')
#end if
if not tiled_structure.has_tmatrix():
self.error('tiling matrix is missing')
#end if
tmatrix = tiled_structure.tmatrix.copy()
tiling = np.diag(tmatrix)
diagonal = np.abs(tmatrix-np.diag(tiling)).sum()==0
if not diagonal:
self.error('non-diagonal tilings are not yet supported by save2qmcpack')
#end if
#end if
#end if
s = '### generated conversion text ###\n'
s += 'from PyscfToQmcpack import savetoqmcpack\n'
self.kpoints = None
self.tiled_kpoints = None
self.tiled_kmap = None
if sys_kpoints is None:
s += "savetoqmcpack({0},{1},'{2}')\n".format(sys_var,mf_var,prefix)
elif tiled_structure is None:
s += 'tiling = [{},{},{}]\n'.format(*tiling)
s += 'for n,kp in enumerate({}):\n'.format(kpts_var)
s += " savetoqmcpack({0},{1},'{2}.twistnum_{{}}'.format(str(n).zfill(3)),kmesh=tiling,kpts=[kp],sp_twist=kp)\n".format(sys_var,mf_var,prefix)
s += "#end for\n"
self.kpoints = sys_kpoints.copy()
else:
s += 'tiling = [{},{},{}]\n'.format(*tiling)
s += "sp_kpoints = {}\n".format(render_array(tiled_kpoints,4))
s += "sp_kmap = {}\n".format(render_array(kmap_array,4))
s += 'for n,kp in enumerate(sp_kpoints):\n'
s += " savetoqmcpack({0},{1},'{2}.twistnum_{{}}'.format(str(n).zfill(3)),kmesh=tiling,kpts={3}[sp_kmap[n]],sp_twist=kp,kmap=sp_kmap[n])\n".format(sys_var,mf_var,prefix,kpts_var)
s += "#end for\n"
self.kpoints = sys_kpoints.copy()
self.tiled_kpoints = tiled_kpoints.copy()
self.tiled_kmap = kmap_array.copy()
#end if
s += '### end generated conversion text ###\n'
self.addendum = '\n'+s+'\n'
#end if
if checkpoint and 'chkfile' not in self.values:
if prefix is None:
self.error('cannot set $chkpoint variable\nplease provide input variable "prefix"')
#end if
chkfile = '{}.chk'.format(prefix)
self.chkfile = chkfile
self.assign(chkfile="'"+chkfile+"'")
#end if
#end def __init__
[docs]
def write_text(self,filepath=None):
text = SimulationInputTemplateDev.write_text(self,filepath)
if self.addendum is not None:
text += self.addendum
#end if
return text
#end def write_text
#end class PyscfInput
#end def generate_pyscf_input