Merge pull request #76 from WardLT/multiple_outcars

Added ability to filter which files a parser can read

Author: maxhutch

dfttopif/drivers.py

@@ -9,21 +9,17 @@
 import json
 
 
-def _add_quality_report(directory, pif, inline=True):
+def _add_quality_report(parser, pif, inline=True):
     import tarfile
 
-    # Use VaspParser to identify OUTCAR and POSCAR files
-    parser = VaspParser(directory)
-
-    # If we do not have an INCAR, we cannot run the quality report
+    # if the parser lacks an INCAR, return None
     if parser.incar is None:
-        print("Unable to generate quality report; directory lacks an INCAR file")
-        return
+        return None
 
     # Create the tar file
     tar = tarfile.open("tmp.tar", "w")
-    tar.add(os.path.join(directory, "OUTCAR"))
-    tar.add(os.path.join(directory, "INCAR"))
+    tar.add(parser.outcar, arcname="OUTCAR")
+    tar.add(parser.incar, arcname="INCAR")
     tar.close()
 
     import requests
@@ -44,7 +40,7 @@ def _add_quality_report(directory, pif, inline=True):
     if inline:
         setattr(pif, "quality_report", report)
     else:
-        report_file = os.path.join(directory, "quality_report.txt")
+        report_file = os.path.join(os.path.dirname(parser.outcar), "quality_report.txt")
         with open(report_file, "w") as f:
             f.write(report)
         if report_file[0:2] == "./":
@@ -59,6 +55,7 @@ def _add_quality_report(directory, pif, inline=True):
 
     return pif
 
+
 def tarfile_to_pif(filename, temp_root_dir='', verbose=0):
     """
     Process a tar file that contains DFT data.
@@ -81,8 +78,8 @@ def tarfile_to_pif(filename, temp_root_dir='', verbose=0):
         for i in os.listdir(temp_dir):
             cur_dir = temp_dir + '/' + i
             if os.path.isdir(cur_dir):
-                return directory_to_pif(cur_dir, verbose)
-        return directory_to_pif(temp_dir, verbose)
+                return directory_to_pif(cur_dir, verbose=verbose)
+        return directory_to_pif(temp_dir, verbose=verbose)
     finally:
         shutil.rmtree(temp_dir)
 
@@ -104,14 +101,13 @@ def archive_to_pif(filename, verbose=0):
     raise Exception('Cannot process file type')
 
 
-def directory_to_pif(directory, verbose=0, quality_report=True, inline=True):
+def files_to_pif(files, verbose=0, quality_report=True, inline=True):
     '''Given a directory that contains output from
     a DFT calculation, parse the data and return
     a pif object
 
     Input:
-        directory - String, path to directory containing
-            DFT results
+        files - [str] list of files from which the parser is allowed to read.
         verbose - int, How much status messages to print
 
     Output:
@@ -123,7 +119,7 @@ def directory_to_pif(directory, verbose=0, quality_report=True, inline=True):
     found_parser = False
     for possible_parser in [PwscfParser, VaspParser]:
         try:
-            parser = possible_parser(directory)
+            parser = possible_parser(files)
             found_parser = True
             break
         except InvalidIngesterException:
@@ -195,12 +191,29 @@ def directory_to_pif(directory, verbose=0, quality_report=True, inline=True):
         chem.properties.append(prop)
 
     # Check to see if we should add the quality report
-    if quality_report and isinstance(parser, VaspParser) :
-        _add_quality_report(directory, chem)
+    if quality_report and isinstance(parser, VaspParser):
+        _add_quality_report(parser, chem)
 
     return chem
 
-def convert(files=[], **kwargs):
+
+def directory_to_pif(directory, **kwargs):
+    """
+    Convert a directory to a pif
+    :param directory: Directory to convert to a pif
+    :param kwargs: any additional keyword arguments. (See `files_to_pif`)
+    :return: the created pif
+    """
+
+    # Get the files
+    files = [os.path.join(directory, f) for f in os.listdir(directory)
+             if os.path.isfile(os.path.join(directory, f))]
+
+    # Run the pif
+    return files_to_pif(files, **kwargs)
+
+
+def convert(files, **kwargs):
     """
     Wrap directory to pif as a dice extension
     :param files: a list of files, which must be non-empty
@@ -213,10 +226,8 @@ def convert(files=[], **kwargs):
 
     if len(files) == 1:
         if os.path.isfile(files[0]):
-            return directory_to_pif(os.path.dirname(files[0]), **kwargs)
+            return files_to_pif(files, **kwargs)
         else:
             return directory_to_pif(files[0], **kwargs)
     else:
-        prefix = os.path.join(".", os.path.commonprefix(files))
-        print("Trying to use prefix {} from {}".format(prefix, os.getcwd()))
-        return directory_to_pif(prefix, **kwargs)
+        return files_to_pif(files, **kwargs)

dfttopif/parsers/abinit.py

@@ -11,14 +11,12 @@ class AbinitParser(DFTParser):
     '''
     _label = None
 
-    def __init__(self, directory):
+    def __init__(self, files):
         # Check whether any file has as name ABINIT in the file in the first two lines
-        super(AbinitParser, self).__init__(directory)
-        files = [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))]
-        is_abinit = False
+        super(AbinitParser, self).__init__(files)
         for f in files:
             try:
-                with open(os.path.join(directory, f), 'r') as fp:
+                with open(f, 'r') as fp:
                     for line in [fp.readline(), fp.readline()]:
                         if "ABINIT" in line:
                             is_abinit = True
@@ -34,31 +32,29 @@ def _get_label(self):
          for this calculation
         '''
         if self._label is None:
-            files = [f for f in os.listdir(self._directory) if os.path.isfile(os.path.join(self._directory, f))] 
             foundfiles = False
-            print(files)
-            for f in files:
+            for f in self._files:
                 if ".files" in f: 
                     foundfiles = True
                     self._label = f.split(".")[0]
-                    fp = open(os.path.join(self._directory, self._label + '.files'), 'r')
-                    line = fp.readline().split()[0]
-                    if line != self._label + ".in":
-                       fp.close()
-                       raise Exception('first line must be label.in') 
-                    line = fp.readline().split()[0]
-                    if line != self._label + ".txt":
-                       fp.close()
-                       raise Exception('second line must be label.txt') 
-                    line = fp.readline().split()[0]
-                    if line != self._label + "i":
-                       fp.close()
-                       raise Exception('third line must be labeli') 
-                    line = fp.readline().split()[0]
-                    if line != self._label + "o":
-                       fp.close()
-                       raise Exception('fourth line must be labelo') 
-                    fp.close()
+                    with open(self._label + '.files', 'r') as fp:
+                        line = fp.readline().split()[0]
+                        if line != self._label + ".in":
+                           fp.close()
+                           raise Exception('first line must be label.in')
+                        line = fp.readline().split()[0]
+                        if line != self._label + ".txt":
+                           fp.close()
+                           raise Exception('second line must be label.txt')
+                        line = fp.readline().split()[0]
+                        if line != self._label + "i":
+                           fp.close()
+                           raise Exception('third line must be labeli')
+                        line = fp.readline().split()[0]
+                        if line != self._label + "o":
+                           fp.close()
+                           raise Exception('fourth line must be labelo')
+                        fp.close()
             if foundfiles:
                 return self._label
             else:

dfttopif/parsers/base.py

@@ -40,22 +40,30 @@ class DFTParser(object):
      have the same name in the pif.
     '''
 
-    _directory = None
-    '''Path to directory containing calculation files'''
-    
     _converged = None
     ''' Whether this calculation has converged '''
 
-    def __init__(self, directory):
-        '''Initialize a parser.
+    def __init__(self, files):
+        '''Initialize a parser by defining the list of files that the parser can read from.
         
         Input:
-            directory - String, path to a directory of output files
+            files - [str], list of files usable by this parser.
         Raises:
-            Par - If parser cannot find needed files
+            InvalidIngesterException - If parser cannot find needed files
         '''
+        self._files = files
+
+    @classmethod
+    def generate_from_directory(cls, directory):
+        """Create a parser by defining which input files it will read from.
 
-        self._directory = directory
+        Input:
+            directory - str, directory to read from
+            files - str, list of files from which to search.
+            """
+        files = [os.path.join(directory, f) for f in os.listdir(directory)
+                 if os.path.isfile(os.path.join(directory, f))]
+        return cls(files)
 
     def get_setting_functions(self):
         '''Get a dictionary containing the names of methods
@@ -101,16 +109,19 @@ def get_result_functions(self):
             'Stresses': 'get_stresses'
         }
 
-    def _call_ase(self, func):
+    def _call_ase(self, func, directory):
         '''Make a call to an ASE function.
+
+        Note: I'm about to kill this function, once I'm done with the overhaul
+        ASE assumes filenames, which we do not want
         
         Handles changing directories
         
         Returns: Result of ASE function
         '''
         # Change directories
         old_path = os.getcwd()
-        os.chdir(self._directory)
+        os.chdir(directory)
 
         # Call function
         try: