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
 
      if filename != None:
         with open(filename) as f:
            self.read(f)
 
   
   def read(self, file):
      line = self.skipEmptyLine(file)
      l = line.split()
      if l[0] != 'fHelmholtz':
         raise Exception('Not valid Thermo file')
      else:
         self.fHelmholtz = float(l[1])
      line = file.readline()
      l = line.split()
      self.pressure = float(l[1])
      line = self.skipEmptyLine(file)
 
      while line != '\n':
         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)
 
      while line != '\n':
         if line == 'polymers:\n':
            self.polymers = []
            self.tableLabel = file.readline()
            line = file.readline()
            while line != '\n':
               l = line.split()
               self.polymers.append(Species(l))
               line = file.readline()
 
         if line == 'solvents:\n':
            self.solvents = []
            line = file.readline()
            line = file.readline()
            while line != '\n':
               l = line.split()
               self.solvents.append(Species(l))
               line = file.readline()
 
         if line == 'cellParams:\n':
            self.cellParams = []
            line = file.readline()
            while line != '\n':
               l = line.split()
               self.cellParams.append(float(l[1]))
               line = file.readline()
         
         line = self.skipEmptyLine(file)
 
         if line == '':
            break
 
   
   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'
 
      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) 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)