Merge remote-tracking branch 'brian/master'

Author: willwheelera

crystalmanager.py

@@ -0,0 +1,295 @@
+from manager_tools import resolve_status, update_attributes
+from crystal import CrystalReader
+from propertiesreader import PropertiesReader
+from autorunner import RunnerPBS
+import os
+import pickle as pkl
+import shutil as sh
+import crystal2qmc
+from autopaths import paths
+
+class CrystalManager:
+  """ Internal class managing process of running a DFT job though crystal.
+  Has authority over file names associated with this task."""
+  def __init__(self,writer,runner,creader=None,name='crystal_run',path=None,
+      preader=None,prunner=None,
+      trylev=False,bundle=False,max_restarts=5):
+    ''' CrystalManager manages the writing of a Crystal input file, it's running, and keeping track of the results.
+    Args:
+      writer (PySCFWriter): writer for input.
+      reader (PySCFReader): to read PySCF output.
+      runner (runner object): to run job.
+      creader (CrystalReader): Reads the crystal results, (None implies use default reader).
+      preader (PropertiesReader): Reads properties results, if any (None implies use default reader).
+      prunner (runner object): run properties if needed (None implies use same runner as crystal).
+      name (str): identifier for this job. This names the files associated with run.
+      trylev (bool): When restarting use LEVSHIFT option to encourage convergence, then do a rerun without LEVSHIFT.
+      bundle (bool): Whether you'll use a bundling tool to run these jobs.
+      max_restarts (int): maximum number of times you'll allow restarting before giving up (and manually intervening).
+    '''
+    # Where to save self.
+    self.name=name
+    self.pickle="%s.pkl"%self.name
+
+    # Ensure path is set up correctly.
+    if path is None:
+      path=os.path.getcwd()
+    if path[-1]!='/': path+='/'
+    self.path=path
+
+    self.logname="%s@%s"%(self.__class__.__name__,self.path+self.name)
+
+    #print(self.logname,": initializing")
+
+    # Handle reader and runner defaults.
+    self.writer=writer
+    if creader is None: self.creader=CrystalReader()
+    else: self.creader=creader
+    if preader is None: self.preader=PropertiesReader()
+    else: self.preader=preader
+    if prunner is None: self.prunner=runner
+    else: self.prunner=prunner
+    if runner is None: self.runner=RunnerPBS()
+    else: self.runner=runner
+
+    # Internal.
+    self.crysinpfn=self.name
+    self.propinpfn=self.name+'.prop'
+    self.crysoutfn=self.crysinpfn+'.o'
+    self.propoutfn=self.propinpfn+'.o'
+    self.restarts=0
+    self._runready=False
+    self.scriptfile=None
+    self.completed=False
+    self.bundle=bundle
+    self.qwfiles={ 
+        'kpoints':[],
+        'basis':'',
+        'jastrow2':'',
+        'orbplot':{},
+        'orb':{},
+        'sys':{},
+        'slater':{}
+      }
+
+    # Smart error detection.
+    self.trylev=trylev
+    self.max_restarts=max_restarts
+    self.savebroy=[]
+    self.lev=False
+
+    # Handle old results if present.
+    if os.path.exists(self.path+self.pickle):
+      #print(self.logname,": rebooting old manager.")
+      old=pkl.load(open(self.path+self.pickle,'rb'))
+      self.recover(old)
+
+    # Update the file.
+    if not os.path.exists(self.path): os.mkdir(self.path)
+    with open(self.path+self.pickle,'wb') as outf:
+      pkl.dump(self,outf)
+
+  #------------------------------------------------
+  def recover(self,other):
+    ''' Recover old class by copying over data. Retain variables from old that may change final answer.'''
+    # Practically speaking, the run will preserve old `take_keys` and allow new changes to `skip_keys`.
+    # This is because you are taking the attributes from the older instance, and copying into the new instance.
+
+    update_attributes(copyto=self,copyfrom=other,
+        skip_keys=['writer','runner','creader','preader','prunner','lev','savebroy',
+                   'path','logname','name',
+                   'trylev','max_restarts','bundle'],
+        take_keys=['restarts','completed','qwfiles'])
+
+    # Update queue settings, but save queue information.
+    update_attributes(copyto=self.runner,copyfrom=other.runner,
+        skip_keys=['queue','walltime','np','nn','jobname'],
+        take_keys=['queueid'])
+    update_attributes(copyto=self.prunner,copyfrom=other.prunner,
+        skip_keys=['queue','walltime','np','nn','jobname'],
+        take_keys=['queueid'])
+
+    update_attributes(copyto=self.creader,copyfrom=other.creader,
+        skip_keys=[],
+        take_keys=['completed','output'])
+
+    update_attributes(copyto=self.preader,copyfrom=other.preader,
+        skip_keys=[],
+        take_keys=['completed','output'])
+
+    updated=update_attributes(copyto=self.writer,copyfrom=other.writer,
+        skip_keys=['maxcycle','edifftol'],
+        take_keys=['completed','modisymm','restart','guess_fort','_elements'])
+    if updated:
+      self.writer.completed=False
+
+  #----------------------------------------
+  def nextstep(self):
+    ''' Determine and perform the next step in the calculation.'''
+    self.recover(pkl.load(open(self.path+self.pickle,'rb')))
+
+    print(self.logname,": next step.")
+    cwd=os.getcwd()
+    os.chdir(self.path)
+
+    # Generate input files.
+    if not self.writer.completed:
+      if self.writer.guess_fort is not None:
+        sh.copy(self.writer.guess_fort,'fort.20')
+      with open(self.crysinpfn,'w') as f:
+        self.writer.write_crys_input(self.crysinpfn)
+      with open(self.propinpfn,'w') as f:
+        self.writer.write_prop_input(self.propinpfn)
+
+    # Check on the CRYSTAL run
+    status=resolve_status(self.runner,self.creader,self.crysoutfn)
+    print(self.logname,": status= %s"%(status))
+
+    if status=="not_started":
+      self.runner.add_command("cp %s INPUT"%self.crysinpfn)
+      self.runner.add_task("%s &> %s"%(paths['Pcrystal'],self.crysoutfn))
+
+    elif status=="ready_for_analysis":
+      #This is where we (eventually) do error correction and resubmits
+      status=self.creader.collect(self.crysoutfn)
+      print(self.logname,": status %s"%status)
+      if status=='killed':
+        if self.restarts >= self.max_restarts:
+          print(self.logname,": restarts exhausted (%d previous restarts). Human intervention required."%self.restarts)
+        else:
+          print(self.logname,": attempting restart (%d previous restarts)."%self.restarts)
+          self.writer.restart=True
+          if self.trylev:
+            print(self.logname,": trying LEVSHIFT.")
+            self.writer.levshift=[10,1] # No mercy.
+            self.savebroy=deepcopy(self.writer.broyden)
+            self.writer.broyden=[]
+            self.lev=True
+          sh.copy(self.crysinpfn,"%d.%s"%(self.restarts,self.crysinpfn))
+          sh.copy(self.crysoutfn,"%d.%s"%(self.restarts,self.crysoutfn))
+          sh.copy('fort.79',"%d.fort.79"%(self.restarts))
+          self.writer.guess_fort='./fort.79'
+          sh.copy(self.writer.guess_fort,'fort.20')
+          self.writer.write_crys_input(self.crysinpfn)
+          sh.copy(self.crysinpfn,'INPUT')
+          self.runner.add_task("%s &> %s"%(paths['Pcrystal'],self.crysoutfn))
+          self.restarts+=1
+    elif status=='done' and self.lev:
+      # We used levshift to converge. Now let's restart to be sure.
+      print("Recovering from LEVSHIFTer.")
+      self.writer.restart=True
+      self.writer.levshift=[]
+      self.creader.completed=False
+      self.lev=False
+      sh.copy(self.crysinpfn,"%d.%s"%(self.restarts,self.crysinpfn))
+      sh.copy(self.crysoutfn,"%d.%s"%(self.restarts,self.crysoutfn))
+      sh.copy('fort.79',"%d.fort.79"%(self.restarts))
+      self.writer.guess_fort='./fort.79'
+      sh.copy(self.writer.guess_fort,'fort.20')
+      self.writer.write_crys_input(self.crysinpfn)
+      sh.copy(self.crysinpfn,'INPUT')
+      self.runner.add_task("%s &> %s"%(paths['Pcrystal'],self.crysoutfn))
+      self.restarts+=1
+
+    # Ready for bundler or else just submit the jobs as needed.
+    if self.bundle:
+      self.scriptfile="%s.run"%self.name
+      self.bundle_ready=self.runner.script(self.scriptfile)
+    else:
+      qsubfile=self.runner.submit(self.path.replace('/','-')+self.name)
+
+    self.completed=self.creader.completed
+
+    # Update the file.
+    with open(self.pickle,'wb') as outf:
+      pkl.dump(self,outf)
+    os.chdir(cwd)
+
+  #----------------------------------------
+  def collect(self):
+    ''' Call the collect routine for readers.'''
+    print(self.logname,": collecting results.")
+    self.creader.collect(self.path+self.crysoutfn)
+
+    # Update the file.
+    with open(self.path+self.pickle,'wb') as outf:
+      pkl.dump(self,outf)
+
+  #------------------------------------------------
+  def script(self,jobname=None):
+    ''' Script execution lines for a bundler to pick up and run.'''
+    if jobname is None: jobname=self.runner.jobname
+    self.scriptfile="%s.run"%jobname
+    self._runready=self.runner.script(self.scriptfile)
+
+  #------------------------------------------------
+  def submit(self,jobname=None):
+    ''' Submit the runner's job to the queue. '''
+    qsubfile=self.runner.submit(jobname)
+    return qsubfile
+
+  #----------------------------------------
+  def to_json(self):
+    raise NotImplementedError
+
+  #----------------------------------------
+  def write_summary(self):
+    self.creader.write_summary()
+    
+  #------------------------------------------------
+  def export_qwalk(self):
+    ''' Export QWalk input files into current directory.'''
+    self.recover(pkl.load(open(self.path+self.pickle,'rb')))
+
+    ready=False
+    if len(self.qwfiles['slater'])==0:
+      self.nextstep()
+
+      if not self.completed:
+        return False
+
+      cwd=os.getcwd()
+      os.chdir(self.path)
+
+      print(self.logname,": %s attempting to generate QWalk files."%self.name)
+
+      # Check on the properties run
+      status=resolve_status(self.prunner,self.preader,self.propoutfn)
+      print(self.logname,": properties status= %s"%(status))
+      if status=='not_started':
+        ready=False
+        self.prunner.add_command("cp %s INPUT"%self.propinpfn)
+        self.prunner.add_task("%s &> %s"%(paths['Pproperties'],self.propoutfn))
+
+        if self.bundle:
+          self.scriptfile="%s.run"%self.name
+          self.bundle_ready=self.prunner.script(self.scriptfile,self.driverfn)
+        else:
+          qsubfile=self.runner.submit(self.path.replace('/','-')+self.name)
+      elif status=='ready_for_analysis':
+        self.preader.collect(self.propoutfn)
+
+      if self.preader.completed:
+        ready=True
+        print(self.logname,": converting crystal to QWalk input now.")
+        self.qwfiles=crystal2qmc.convert_crystal(base=self.name,propoutfn=self.propoutfn)
+      else:
+        ready=False
+        print(self.logname,": conversion postponed because properties is not finished.")
+
+      os.chdir(cwd)
+    else:
+      ready=True
+
+    with open(self.path+self.pickle,'wb') as outf:
+      pkl.dump(self,outf)
+
+    return ready
+    
+  #----------------------------------------
+  def status(self):
+    if self.completed:
+      return 'ok'
+    else:
+      return 'not_finished'
+

intro/00-introduction.ipynb

@@ -47,7 +47,8 @@
     "from autorunner import PySCFRunnerLocal,RunnerLocal\n",
     "from trialfunc import SlaterJastrow\n",
     "from variance import VarianceWriter, VarianceReader\n",
-    "from manager import PySCFManager, QWalkManager\n",
+    "from pyscfmanager import PySCFManager\n",
+    "from qwalkmanager import QWalkManager\n",
     "\n",
     "path='00-scratch'\n",
     "\n",

intro/00-introduction.py

@@ -32,7 +32,8 @@
 from autorunner import PySCFRunnerLocal,RunnerLocal
 from trialfunc import SlaterJastrow
 from variance import VarianceWriter, VarianceReader
-from manager import PySCFManager, QWalkManager
+from pyscfmanager import PySCFManager
+from qwalkmanager import QWalkManager
 
 path='00-scratch'
 

intro/04-job_management.ipynb

@@ -82,7 +82,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from manager import PySCFManager\n",
+    "from pyscfmanager import PySCFManager\n",
     "pyscf_manager = PySCFManager(\n",
     "    path='04-scratch',\n",
     "    name='h2_pbe',\n",

intro/04-job_management.py

@@ -59,7 +59,7 @@
 # In[ ]:
 
 
-from manager import PySCFManager
+from pyscfmanager import PySCFManager
 pyscf_manager = PySCFManager(
     path='04-scratch',
     name='h2_pbe',

intro/05-qwalk.ipynb

@@ -39,7 +39,7 @@
    "source": [
     "from autopyscf import PySCFWriter,PySCFReader\n",
     "from autorunner import PySCFRunnerLocal\n",
-    "from manager import PySCFManager\n",
+    "from pyscfmanager import PySCFManager\n",
     "\n",
     "h2='\\n'.join([\n",
     "    'H 0.0 0.0 0.0 ',\n",
@@ -95,7 +95,7 @@
    "source": [
     "from variance import VarianceWriter,VarianceReader\n",
     "from autorunner import RunnerLocal\n",
-    "from manager import QWalkManager\n",
+    "from qwalkmanager import QWalkManager\n",
     "from imp import reload\n",
     "import time\n",
     "\n",

intro/05-qwalk.py

@@ -24,7 +24,7 @@
 
 from autopyscf import PySCFWriter,PySCFReader
 from autorunner import PySCFRunnerLocal
-from manager import PySCFManager
+from pyscfmanager import PySCFManager
 
 h2='\n'.join([
     'H 0.0 0.0 0.0 ',
@@ -64,7 +64,7 @@
 
 from variance import VarianceWriter,VarianceReader
 from autorunner import RunnerLocal
-from manager import QWalkManager
+from qwalkmanager import QWalkManager
 from imp import reload
 import time