Merge pull request #2858 from jtkrogel/nx_batched_spindens

Nexus: support batched spin density estimator

nexus/bin/qdens

@@ -60,6 +60,7 @@ try:
     qmcpack_input = import_nexus_module('qmcpack_input')
     QmcpackInput    = qmcpack_input.QmcpackInput
     spindensity_xml = qmcpack_input.spindensity
+    spindensity_new_xml = qmcpack_input.spindensity_new  # temporary
     del qmcpack_input
 except:
     from generic import obj
@@ -71,6 +72,7 @@ except:
     from numerics import simplestats,simstats
     from qmcpack_input import QmcpackInput
     from qmcpack_input import spindensity as spindensity_xml
+    from qmcpack_input import spindensity_new as spindensity_new_xml # temporary
 #end try
 
 
@@ -901,29 +903,55 @@ class QMCDensityProcessor(QDBase):
             structure = s
             cell = s.axes.copy()
             qi.pluralize()
-            est = qi.get('estimators')
-            if est is not None and len(est)==1 and 'LocalEnergy' in est:
-                ham = qi.get('hamiltonian')
+            # collect every bunch of estimators strewn throughout the input file
+            est_sources = []
+            ham = qi.get('hamiltonian')
+            if ham is not None:
                 if 'estimators' in ham:
-                    est = ham.estimators
-                elif len(ham)==1:
-                    est = ham.first().get('estimators')
+                    est_sources.append(ham.estimators)
+                elif len(ham)>0:
+                    ham = ham.first() # hamiltonian collection
+                    if 'estimators' in ham:
+                        est_sources.append(ham.estimators)
+                    #end if
                 #end if
             #end if
-            if est is not None:
-                opt.grids = obj()
+            calcs = qi.get('calculations')
+            if calcs is not None:
+                for series in sorted(calcs.keys()):
+                    qmc = calcs[series]
+                    if 'estimators' in qmc:
+                        est_sources.append(qmc.estimators)
+                    #end if
+                #end for
+            #end if
+            # collect input information from the first spin density instance with a particular name
+            #   assume all other instances that share this name have identical inputs
+            grids = obj()
+            for est in est_sources:
                 for name,xml in est.items():
-                    if isinstance(xml,spindensity_xml):
+                    # the if statement below is a nasty hack to temporarily support 
+                    # the name mismatch between the input file and stat.h5 currently 
+                    # enforced by the (new) batched spin denisity class in qmcpack
+                    if isinstance(xml,spindensity_new_xml):
+                        name = 'SpinDensity' # override all user input, just like qmcpack
+                    #end if
+                    if name not in grids and isinstance(xml,(spindensity_xml,spindensity_new_xml)):
                         sd = xml
                         if 'grid' in sd:
-                            opt.grids[name] = sd.grid
+                            grids[name] = sd.grid
                         elif 'dr' in sd:
-                            opt.grids[name] = get_grid(s,sd.dr)
+                            grids[name] = get_grid(s,sd.dr)
                         else:
                             self.error('could not identify grid data for spin density named "{0}"\bin QMCPACK input file: {1}'.format(name,opt.input))
                         #end if
                     #end if
                 #end for
+            #end for
+            if len(grids)>0:
+                opt.grids = grids
+            else:
+                self.error('Could not find any spin density estimators in the input file provided.\nInput file provided: {}'.format(opt.input))
             #end if
         #end if
 

nexus/lib/qmcpack_input.py

@@ -2223,6 +2223,14 @@ class spindensity(QIxml):
     identifier  = 'name'
 #end class spindensity
 
+class spindensity_new(QIxml): # temporary
+    tag = 'estimator'
+    attributes  = ['type','name','report','save_memory']
+    parameters  = ['dr','grid','cell','center','corner','voronoi','test_moves']
+    write_types = obj(report=yesno,save_memory=yesno)
+    identifier  = 'name'
+#end class spindensity_new
+
 class structurefactor(QIxml):
     tag = 'estimator'
     attributes  = ['type','name','report']
@@ -2323,6 +2331,7 @@ estimator = QIxmlFactory(
                  nearestneighbors    = nearestneighbors,
                  dm1b                = dm1b,
                  spindensity         = spindensity,
+                 spindensity_new     = spindensity_new, # temporary
                  structurefactor     = structurefactor,
                  force               = force,
                  forwardwalking      = forwardwalking,
@@ -2520,6 +2529,15 @@ class rmc(QIxml):
     write_types = obj(collect=yesno)
 #end class rmc
 
+class vmc_batch(QIxml):
+    collection_id = 'qmc'
+    tag = 'qmc'
+    attributes = ['method','move']
+    elements   = ['estimator']
+    parameters = ['walkers','warmupsteps','blocks','steps','substeps','timestep','usedrift']
+    write_types = obj(usedrift=yesno)
+#end class vmc_batch
+
 class wftest(QIxml):
     collection_id = 'qmc'
     tag = 'qmc'
@@ -2539,7 +2557,7 @@ class setparams(QIxml):
 
 qmc = QIxmlFactory(
     name = 'qmc',
-    types   = dict(linear=linear,cslinear=cslinear,vmc=vmc,dmc=dmc,loop=loop,optimize=optimize_qmc,wftest=wftest,rmc=rmc,setparams=setparams),
+    types   = dict(linear=linear,cslinear=cslinear,vmc=vmc,dmc=dmc,loop=loop,optimize=optimize_qmc,wftest=wftest,rmc=rmc,setparams=setparams,vmc_batch=vmc_batch),
     typekey = 'method',
     default = 'loop'
     )
@@ -2596,10 +2614,11 @@ classes = [   #standard classes
     localenergy,energydensity,spacegrid,origin,axis,wavefunction,
     determinantset,slaterdeterminant,basisset,grid,determinant,occupation,
     jastrow1,jastrow2,jastrow3,
-    correlation,coefficients,loop,linear,cslinear,vmc,dmc,
+    correlation,coefficients,loop,linear,cslinear,vmc,dmc,vmc_batch,
     atomicbasisset,basisgroup,init,var,traces,scalar_traces,particle_traces,array_traces,
     reference_points,nearestneighbors,neighbor_trace,dm1b,
     coefficient,radfunc,spindensity,structurefactor,
+    spindensity_new, # temporary
     sposet,bspline_builder,composite_builder,heg_builder,include,
     multideterminant,detlist,ci,mcwalkerset,csf,det,
     optimize,cg_optimizer,flex_optimizer,optimize_qmc,wftest,kspace_jastrow,
@@ -2816,6 +2835,9 @@ pressure.defaults.set(
 momentum.defaults.set(
     type='momentum'
     )
+spindensity_new.defaults.set( # temporary
+    type='spindensity_new',name='SpinDensityNew'
+    )
 
 
 linear.defaults.set(
@@ -2912,6 +2934,9 @@ dmc.defaults.set(
     #nonlocalmoves = True,
     #estimators = classcollection(localenergy)
     )
+vmc_batch.defaults.set(
+    method='vmc_batch',move='pbyp',
+    )