# Author: Yubo "Paul" Yang
# Email: yubo.paul.yang@gmail.com
import os
import subprocess as sp
from lxml import etree
from qharv.reel import mole
from qharv.seed import xml, xml_examples
# =============== level 0: build input from scratch ===============
[docs]def assemble_project(nodel, name='qmc', series=0):
""" assemble QMCPACK input using a list of xml nodes
usually nodel=[qmcsystem, qmc]
Args:
nodel (list): a list of xml node (lxml.Element)
name (str, optional): project name, default 'qmc'
"""
qsim = etree.Element('simulation')
proj = xml.make_node('project', {'id': name, 'series': str(series)})
for node in [proj]+nodel:
qsim.append(node)
doc = etree.ElementTree(qsim)
return doc
[docs]def simulationcell_from_axes(axes, bconds='p p p', rckc=15.):
""" construct the <simulationcell> xml element from axes
Args:
axes (np.array): lattice vectors
bconds (str, optional): boundary conditions in x,y,z directions.
p for periodic, n for non-periodic, default to 'p p p'
rckc: long-range cutoff paramter rc*kc, default to 15
Return:
etree.Element: representing <simulationcell>
"""
def pad_line(line): # allow content to be selected by double clicked
return ' ' + line + ' '
# write primitive lattice vectors
lat_node = etree.Element('parameter', attrib={
'name': 'lattice',
'units': 'bohr'
})
lat_node.text = xml.arr2text(axes)
# write boundary conditions
bconds_node = etree.Element('parameter', {'name': 'bconds'})
bconds_node.text = pad_line(bconds)
# write long-range cutoff parameter
lr_node = etree.Element('parameter', {'name': 'LR_dim_cutoff'})
lr_node.text = pad_line(str(rckc))
# build <simulationcell>
sc_node = etree.Element('simulationcell')
sc_node.append(lat_node)
sc_node.append(bconds_node)
sc_node.append(lr_node)
return sc_node
[docs]def particle_group_from_pos(pos, name, charge, **kwargs):
""" construct a <group> in the <particleset> xml element
Args:
pos (np.array): positions, shape (nptcl, ndim)
name (str): name of particle group
charge (float): the amount of charge of this particle species
Return:
etree.Element: <group> including <attrib name="positions">
"""
if 'charge' in kwargs:
msg = 'keyword charge %f will be overwriten'
msg += ' by %f' % (kwargs['charge'], charge)
raise RuntimeError(msg)
kwargs['charge'] = charge
group = xml.etree.Element('group', dict(
name = name,
size = str(len(pos)),
))
for key, val in kwargs.items():
xml.set_param(group, key, str(val), new=True)
pos_attrib = xml.etree.Element('attrib', dict(
name = 'position',
datatype = 'posArray',
condition = str(0)
))
pos_attrib.text = xml.arr2text(pos)
group.append(pos_attrib)
return group
[docs]def ud_electrons(nup, ndown):
""" construct the <particleset name="e"> xml element for electrons
Args:
nup (int): number of up electrons
ndown (int): number of down electrons
Return:
etree.Element: representing <particleset name="e">
"""
epset = etree.Element('particleset', {'name': 'e', 'random': 'yes'})
up_group = etree.Element('group', {
'name': 'u',
'size': str(nup),
'mass': '1.0'
})
dn_group = etree.Element('group', {
'name': 'd',
'size': str(ndown),
'mass': '1.0'
})
for egroup in [up_group, dn_group]:
xml.set_param(egroup, 'charge', ' -1 ', new=True)
epset.append(egroup)
return epset
[docs]def all_electron_hamiltonian(elec_name='e', ion_name='ion0'):
ee = xml.make_node('pairpot', {
'type': 'coulomb', 'name': 'ElecElec',
'source': elec_name, 'target': elec_name
})
ei = xml.make_node('pairpot', {
'type': 'coulomb', 'name': 'ElecIon',
'source': ion_name, 'target': elec_name
})
ii = xml.make_node('pairpot', {
'type': 'coulomb', 'name': 'IonIon',
'source': ion_name, 'target': ion_name
})
ham = xml.make_node('hamiltonian')
xml.append(ham, [ee, ei, ii])
return ham
[docs]def bspline_qmcsystem(fh5, tmat=None):
"""Create Slater-Jastrow system input from pw2qmcpack.x h5 file
Args:
fh5 (str): path to wf h5 file
tmat (np.array): tile matrix
Return:
qsys (etree.Element): <qmcsystem>
"""
import numpy as np
from qharv.seed import wf_h5
ndim0 = 3 # !!!! hard-code for three dimensions
fp = wf_h5.read(fh5)
axes, elem, charge_map, pos = wf_h5.axes_elem_charges_pos(fp)
nelecs = wf_h5.get(fp, 'nelecs')
fp.close()
natom, ndim = pos.shape
assert ndim == ndim0
nup, ndn = nelecs
if nup != ndn: # hard-code for unpolarized for now
raise RuntimeError('nup != ndn')
if tmat is None: # use primitive cell by default
tmat = np.eye(3, dtype=int)
else: # tile supercell
from qharv.inspect.axes_elem_pos import ase_tile
axes, elem, pos = ase_tile(axes, elem, pos, tmat)
spoup = spodn = 'spo_ud'
psi_name = 'psi0'
ion_name = 'ion0'
nodes = []
# simulationcell
sc_node = simulationcell_from_axes(axes)
nodes.append(sc_node)
# particlesets
ions = xml.make_node('particleset', {'name': ion_name})
for name, charge in charge_map.items():
sel = elem == name
ion_grp = particle_group_from_pos(pos[sel], name, charge)
ions.append(ion_grp)
nodes.append(ions)
elecs = ud_electrons(*nelecs)
nodes.append(elecs)
# sposet and builder
tmat_str = ('%d ' * 9) % tuple(tmat.ravel())
sb = xml.make_node('sposet_builder', {
'type': 'bspline',
'href': fh5,
'tilematrix': tmat_str,
'twistnum': '0',
'source': ion_name, # "Einspline needs the source particleset"
})
spo = xml.make_node('sposet', {
'type': 'bspline',
'name': spoup,
'size': str(nup),
'spindataset': '0'}
)
sb.append(spo)
# determinantset
updet = xml.make_node('determinant', {
'id': 'updet',
'size': str(nup),
'sposet': spoup
})
dndet = xml.make_node('determinant', {
'id': 'dndet',
'size': str(ndn),
'sposet': spodn
})
sdet = xml.make_node('slaterdeterminant')
xml.append(sdet, [updet, dndet])
dset = xml.make_node('determinantset')
dset.append(sdet)
# wave function
wf = xml.make_node('wavefunction', {'name': psi_name, 'target': 'e'})
xml.append(wf, [sb, dset])
nodes.append(wf)
# hailtonian
ham = all_electron_hamiltonian()
nodes.append(ham)
qsys = xml.make_node('qmcsystem')
xml.append(qsys, nodes)
return qsys
# ================== level 1: use existing input ===================
[docs]def expand_twists(example_in_xml, twist_list, calc_dir, force=False):
""" expand example input xml to all twists in twist_list
Naming convention of new inputs:
[prefix].g001.twistnum_[itwist].in.xml
Args:
example_in_xml (str): example QMCPACK input xml file
twist_list (list): a list of twist indices
calc_dir (str): folder to output new inputs
Return:
None
Examples:
>>> expand_twists('./vmc.in.xml',range(64),'.')
>>> expand_twists('./ref/vmc.in.xml',[0,15,17],'./new')
"""
doc = xml.read(example_in_xml)
prefix = doc.find('.//project').get('id')
fname_fmt = '{prefix:s}.{gt:s}.twistnum_{itwist:d}.in.xml'
bundle_text = ''
bundle_in = os.path.join(calc_dir, '%s.in' % prefix)
for itwist in twist_list:
# change twist number
bb = doc.find('.//sposet_builder[@type="bspline"]')
if bb is None: # assume old-style input
bb = doc.find('.//determinantset[@type="bspline"]')
bb.set('twistnum', str(itwist))
# construct file name
gt = 'g' + str(itwist).zfill(3)
fname = fname_fmt.format(
prefix = prefix,
gt = gt,
itwist = itwist
)
floc = os.path.join(calc_dir, fname)
bundle_text += '%s\n' % fname
if not force: # check if file exists
if os.path.isfile(floc):
raise RuntimeError('force to overwrite %s' % floc)
xml.write(floc, doc)
with open(bundle_in, 'w') as f:
f.write(bundle_text)
[docs]def bundle_twists(calc_dir, fregex='*twistnum_*.in.xml'):
""" bundle all twist inputs
quick and dirty: `cd $calc_dir; ls > prefix.in`, then edit prefix.in
Args:
calc_dir (str): calculation directory
fregex (str, optional): regular expression for all twists
Return:
str: bundled input text
"""
flist = mole.files_with_regex(fregex, calc_dir, maxdepth=1)
flist.sort()
text = ''
for floc in flist:
fname = os.path.basename(floc)
text += fname + '\n'
return text
[docs]def disperse(ginp_loc, calc_dir, execute=False, overwrite=False):
""" disperse inputs bundled up in a grouped input
Args:
ginp_loc (str): location of grouped input e.g. ../runs/dmc/qmc.in
calc_dir (str): folder to output new inputs e.g. dmc1
execute (bool,optional): perform file I/O, default is False i.e. a dry run
overwrite (bool,optional): overwrite existing files, default is False
Returns:
list: a list of new inputs
"""
# path0 is the folder containing the current grouped input
path0 = os.path.dirname(ginp_loc)
calc_dir0 = os.path.basename(path0)
# path is the folder to contain the dispersed inputs
path = os.path.join(os.path.dirname(path0), calc_dir)
if execute: # make folder if not there
if not os.path.isdir(path):
sp.check_call(['mkdir', path])
# for each input in grouped input file, add group text (gt) to project id
# if execute, write input in given folder
flist = []
with open(ginp_loc, 'r') as f:
ig = 0
for line in f:
# construct source and target input paths
infile = line.strip('\n')
floc0 = os.path.join(path0, infile)
if not os.path.isfile(floc0):
raise RuntimeError('%s not found' % floc0)
floc = os.path.join(path, infile)
if os.path.isfile(floc) and (not overwrite) and execute:
raise RuntimeError('%s exists; delete or overwrite ' % floc)
flist.append(floc)
# modify prefix
gt = 'g'+str(ig).zfill(3)
doc = xml.read(floc0)
pnode = doc.find('.//project')
prefix0 = pnode.get('id')
prefix = '.'.join([prefix0, gt])
pnode.set('id', prefix)
if execute:
xml.write(floc, doc)
ig += 1
return flist
[docs]def set_norb(doc, norb):
""" occupy the lowest norb Kohn-Sham orbitals
Args:
doc (etree.Element): must contain <particleset>, <sposet>, and <det...set>
norb (int): number of orbitals to occupy
Return:
None
Effect:
doc is modified
"""
epset = doc.find('.//particleset[@name="e"]')
for group in epset.findall('.//group'): # 'u' and 'd'
group.set('size', str(norb))
sposet = doc.find('.//sposet[@name="spo_ud"]')
sposet.set('size', str(norb))
detset = doc.find('.//determinantset')
for det in detset.findall('.//determinant'):
det.set('size', str(norb))
[docs]def set_gc_occ(norbl, calc_dir, fregex_fmt='*twistnum_{itwist:d}.in.xml'):
""" edit twist inputs in calc_dir according to occupation vector norbl
Args:
norbl (list): number of occupied orbitals at each twist
calc_dir (str): location of twist inputs
"""
nocc = len(norbl)
wild_fregex = fregex_fmt.replace('{itwist:d}', '*')
flist = mole.files_with_regex(wild_fregex, calc_dir)
ntwist = len(flist)
if nocc != ntwist:
raise RuntimeError('%d occupations given for %d twists' % (nocc, ntwist))
for itwist in range(ntwist):
norb = int(norbl[itwist])
fregex = fregex_fmt.format(itwist=itwist)
fxml = mole.find(fregex, calc_dir)
doc = xml.read(fxml)
set_norb(doc, norb)
xml.write(fxml, doc)
[docs]def set_nwalker(doc, nwalker):
""" set the number of walkers to use in DMC
Args:
doc (lxml.Element): xml node containing <qmc>
nwalker (int): number of walkers
"""
nodes = doc.findall('.//qmc[@method="vmc"]')
for node in nodes:
xml.set_param(node, 'samples', str(nwalker))
nodes = doc.findall('.//qmc[@method="dmc"]')
for node in nodes:
xml.set_param(node, 'targetwalkers', str(nwalker))