output.py

"""! Module of parsers for PSCF output file formats. """
 
import pscfpp.param as param
import os
 
# Contains classes Thermo, Species, State, and Sweep
 

class Thermo:
 
   
   def __init__(self, filename=None):
      self.fHelmholtz = None
      self.pressure = None
      self.fIdeal = None
      self.fInter = None
      self.fExt = None
      self.polymers = None
      self.solvents = None
      self.cellParams = None
      self.tableLabel = None
      self.homogeneous = None
      self.homogeneousSpecies = None
      self.stress = None
      self.environmentStress = None
 
      if filename != None:
         with open(filename) as f:
            self.read(f)
 
   
   def read(self, file):
      line = self.skipEmptyLine(file)
 
      # Read fHelmholtz and pressure (required)
      if line.startswith('fHelmholtz'):
         self.fHelmholtz = float(line.split()[-1])
      else:
         raise Exception('Not valid Thermo file. No fHelmholtz found.')
      line = file.readline()
      if line.startswith('pressure'):
         self.pressure = float(line.split()[-1])
      else:
         raise Exception('Not valid Thermo file. No pressure found.')
      line = self.skipEmptyLine(file)
 
      # Now, read all other blocks of thermo file in any order
      hasPolymersBlock = False
      while line != '':
 
         # Read free energy contributions (optional)
         if line.startswith('fIdeal'):
            while line.strip() != '':
               l = line.split()
               if l[0] == 'fIdeal':
                  self.fIdeal = float(l[-1])
               if l[0] == 'fInter':
                  self.fInter = float(l[-1])
               if l[0] == 'fExt':
                  self.fExt = float(l[-1])
               line = file.readline()
            line = self.skipEmptyLine(file)
 
         # Read polymers block
         elif line.lower().startswith('polymers:'):
            self.polymers = []
            self.tableLabel = file.readline()
            line = file.readline()
            while line.strip() != '':
               l = line.split()
               self.polymers.append(Species(l))
               line = file.readline()
            line = self.skipEmptyLine(file)
            hasPolymersBlock = True
 
         # Read solvents block (optional)
         elif line.lower().startswith('solvents:'):
            self.solvents = []
            line = file.readline()
            line = file.readline()
            while line.strip() != '':
               l = line.split()
               self.solvents.append(Species(l))
               line = file.readline()
            line = self.skipEmptyLine(file)
 
         # Read lattice parameters (optional)
         elif line.startswith('cellParams:'):
            self.cellParams = []
            line = file.readline()
            while line.strip() != '':
               l = line.split()
               self.cellParams.append(float(l[1]))
               line = file.readline()
            line = self.skipEmptyLine(file)
 
         # Read homogeneous block (optional, r1d only)
         elif line.startswith('f (homo)'):
            self.homogeneous = {}
            self.homogeneousSpecies = []
 
            while line.strip() != '':
               self.homogeneous[line[:11].strip()] = float(line[12:31])
               line = file.readline()
            
            line = self.skipEmptyLine(file)
            if line.lower().startswith('species:'):
               line = file.readline()
               speciesParams = line.split()
               line = file.readline()
               while line.strip() != '':
                  l = line.split()
                  self.homogeneousSpecies.append({})
                  for i in range(4):
                     prm = speciesParams[i]
                     if i == 0:
                        self.homogeneousSpecies[-1][prm] = int(l[i])
                     else:
                        self.homogeneousSpecies[-1][prm] = float(l[i])
                  line = file.readline()
            line = self.skipEmptyLine(file)
      
         # Read stress block (optional)
         elif line.lower().startswith('stress:'):
            self.stress = []
            line = file.readline()
            while line.strip() != '':
               l = line.split()
               self.stress.append(float(l[1]))
               line = file.readline()
            line = self.skipEmptyLine(file)
         
         # Read environment stress block (optional)
         elif line.lower().startswith('environment-modified stress:'):
            self.environmentStress = []
            line = file.readline()
            while line.strip() != '':
               l = line.split()
               self.environmentStress.append(float(l[1]))
               line = file.readline()
            line = self.skipEmptyLine(file)
 
      if not hasPolymersBlock:
         raise Exception('Not valid Thermo file. No polymers block found.')
 
   
   def skipEmptyLine(self, file):
      line = file.readline()
      while line == '\n':
         line = file.readline()
      return line
 
   
   def write(self, filename):
      with open(filename, 'w') as f:
         f.write(self.__str__())
 
   
   def __str__(self):
      s = ''
      s += 'fHelmholtz'
      v = f'{self.fHelmholtz:.11e}'
      s += f'{v:>22}\n'
      s += 'pressure  '
      v = f'{self.pressure:.11e}'
      s += f'{v:>22}\n'
      s += '\n'
 
      if (self.fIdeal != None) or (self.fInter != None) or (self.fExt != None):
         if self.fIdeal != None:
            s += 'fIdeal    '
            v = f'{self.fIdeal:.11e}'
            s += f'{v:>22}\n'
         if self.fInter != None:
            s += 'fInter    '
            v = f'{self.fInter:.11e}'
            s += f'{v:>22}\n'
         if self.fExt != None:
            s += 'fExt      '
            v = f'{self.fExt:.11e}'
            s += f'{v:>22}\n'
         s += '\n'
 
      s += 'polymers:\n'
      s += self.tableLabel
      for i in range(len(self.polymers)):
         p = f'{self.polymers[i].phi:.11e}'
         m = f'{self.polymers[i].mu:.11e}'
         s += f'{i:>5}{p:>20}{m:>20}\n'
      s += '\n'
 
      if self.solvents != None:
         s += 'solvents:\n'
         s += self.tableLabel
         for i in range(len(self.solvents)):
            p = f'{self.solvents[i].phi:.11e}'
            m = f'{self.solvents[i].mu:.11e}'
            s += f'{i:>5}{p:>20}{m:>20}\n'
         s += '\n'
 
      if self.cellParams != None:
         s += 'cellParams:\n'
         for i in range(len(self.cellParams)):
            v = f'{self.cellParams[i]:.11e}'
            s += f'{i:>5}{v:>20}\n'
         s += '\n'
 
      if self.homogeneous:
         counter = 0
         for key, val in self.homogeneous.items():
            s += f'{key:<10} = {val:18.11e}'
            if key == "V(tot)/v":
               s += "\n"
            elif key.lower() == key: # key is all lowercase
               s += "   [per monomer volume]\n"
            else:
               s += "   [total]\n"
            counter += 1
         s += '\n'
 
      if self.homogeneousSpecies:
         s += 'Species:\n'
         s += '    i      mu                  phi(homo)           deltaV\n'
         for sp in self.homogeneousSpecies:
            s += f'{sp["i"]:>5}'
            s += f'{sp["mu"]:>20.11e}'
            s += f'{sp["phi(homo)"]:>20.11e}'
            s += f'{sp["deltaV"]:>20.11e}\n'
         s += '\n'
      
      if self.stress != None:
         s += 'stress:\n'
         for i in range(len(self.stress)):
            v = f'{self.stress[i]:.11e}'
            s += f'{i:>5}{v:>20}\n'
         s += '\n'
      
      if self.environmentStress != None:
         s += 'environment-modified stress:\n'
         for i in range(len(self.environmentStress)):
            v = f'{self.environmentStress[i]:.11e}'
            s += f'{i:>5}{v:>20}\n'
         s += '\n'
 
      return s
 

class Species:
   
   
   def __init__(self, l):
      if len(l) == 3:
         self.phi = float(l[1])
         self.mu = float(l[2])
      if len(l) == 2:
         self.phi = float(l[0])
         self.mu = float(l[1])
 
 
 

class State:
 
   
   def __init__(self, filename):
      with open(filename) as f:
         firstline = f.readline()
         fl = firstline.split()
         if fl[0] != 'System{':
            raise Exception('Not valid State file')
         else:
            self.param = param.Composite(f, 'System')
            self.thermo = Thermo()
            self.thermo.read(f)
 
   
   def __str__(self):
      out = self.param.__str__() + '\n' + self.thermo.__str__()
      return out
 
   
   def write(self, filename):
      with open(filename, "w") as f:
         f.write(self.__str__())
 
 

class Sweep:
 
   
   def __init__(self, prefix):
      self.sweep = []
      num = 0
      filename = prefix + str(num) + '.stt'
      while os.path.isfile(filename) == True:
         s = State(filename)
         self.sweep.append(s)
         num += 1
         filename = prefix + str(num) + '.stt'
 
   
   def summary(self, vars, index = False):
      n = len(vars)
 
      summary = []
 
      for i in range(0, len(self.sweep)):
         if index == True:
            s = [i]
         else:
            s = []
         for j in range(0, n):
            a = self.sweep[i]
            string = 'a.' + vars[j]
            try:
               val = eval(string)
            except RecursionError:
               raise Exception('Wrong command or values do not exist')
            s.append(val)
         summary.append(s)
 
      return summary
 
 
   
   def summaryString(self, vars):
      n = len(vars)
 
      summary = []
 
      for i in range(0, len(self.sweep)):
         s = [i]
         for j in range(0, n):
            a = self.sweep[i]
            string = 'a.' + vars[j]
            try:
               val = eval(string)
            except RecursionError:
               raise Exception('Wrong command or values do not exist')
            s.append(val)
         summary.append(s)
 
      nl = [4]
      nameList = ['step']
      for i in range(0, n):
         index = vars[i].rfind('.')
         name = vars[i][index+1:]
         nameList.append(name)
         nl.append(len(name))
 
      valType = []
      for i in range(0, len(summary)):
         valType.append([])
         for j in range(0, len(summary[0])):
            valType[i].append(type(summary[i][j]))
 
      for i in range(0, len(summary)):
         for j in range(0, len(summary[0])):
            length = len(str(summary[i][j]))
            if (valType[i][j] == str) and (length > nl[j]):
               nl[j] = length
            if (valType[i][j] == float) and (13 > nl[j]):
               nl[j] = 13
            if (valType[i][j] == int) and (length > nl[j]):
               nl[j] = length
 
 
      summaryString = ' '
      for i in range(0, len(nameList)):
         stringFormat = '{:>' + str(nl[i]) + 's}'
         if i != len(nameList)-1:
            summaryString += stringFormat.format(nameList[i]) + '  '
         else:
            summaryString += stringFormat.format(nameList[i]) + '\n '
 
      for i in range(0, len(summary)):
         for j in range(0, len(summary[0])):
            if valType[i][j] == int:
               stringFormat = '{:>' + str(nl[j]) + '}'
               summaryString += stringFormat.format(summary[i][j])
            if valType[i][j] == float:
               stringFormat = '{:.7e}'
               val = stringFormat.format(summary[i][j])
               summaryString += val
            if valType[i][j] == str:
               stringFormat = '{:>' + str(nl[j]) + 's}'
               summaryString += stringFormat.format(summary[i][j])
            if j == len(summary[0])-1:
               summaryString += '\n '
            else:
               summaryString += '  '
 
      return summaryString
 
   
   def __getitem__(self, key):
      return self.sweep[key]
 
   
   def __len__(self):
      return len(self.sweep)