analyze_fesp_init.py

#!/usr/bin/env python3
'''
Created by pkawak
Tue Apr  9 16:17:09 EDT 2024

'''

import pandas as pd
import os
import matplotlib.pyplot as plt
import numpy as np
import sys

class Step:
  # Step is a class; each object instance has a number of variables and methods that assist in defining a step in our initialization code
  df            = [] # dataframe to hold thermo output
  fix_strs      = [] # strings with fixes active during thermo output
  fix_ids       = [] # ids of the fixes active during thermo output
  log_file      = "" # log file name with thermo output
  step_line_num = 0  # first line number of thermo output in log_file
  loop_line_num = 0  # last line number of thermo output in log_file

  def preprocess_df(self):
    # preprocesses dataframe by adding necessary columns from preexisting ones
    self.df["P_pol"] = 1/3*(self.df["f_avgPxx_all[1]"]+self.df["f_avgPxx_all[2]"]+self.df["f_avgPxx_all[3]"])
    self.df["P_fil"] = 1/3*(self.df["f_avgPxx_all[7]"]+self.df["f_avgPxx_all[8]"]+self.df["f_avgPxx_all[9]"])
    self.df["P_all"] = self.df["Press"]

  def get_df(self):
    # reads in the dataframe and preprocesses
    self.df = pd.read_csv(self.log_file, index_col=False, delim_whitespace=True, skiprows=self.step_line_num-1, nrows=self.loop_line_num-self.step_line_num-1)
    if self.df.iloc[-1][0] == "Fix":
      self.loop_line_num -= 1
      self.df = pd.read_csv(self.log_file, index_col=False, delim_whitespace=True, skiprows=self.step_line_num-1, nrows=self.loop_line_num-self.step_line_num-1)

    self.preprocess_df()

  def get_fixes(self):
    # goes through log_file and identifies fixes active during thermo output
    with open(self.log_file) as f:
      # read all lines up to the "Step" line in this current segment of file
      lines = f.readlines()[:self.step_line_num]

      # this next part filters all non-fix lines, time averaging fixes, lines with variables, and momentum fixes and also strips all whitespace and converts it into a single space
      lines = [ ' '.join(l.split()) for l in lines if "fix" in l and "ave/" not in l and "${" not in l and "momentum" not in l and "style" not in l or "minimize" in l ]
      lines = [ l.split('#')[0] for l in lines ]

      # keep only lines starting with fix and minimize
      self.fix_strs  = [ l for l in lines if l.startswith("fix") or l.startswith("minimize") ]
      # get ids (2nd word in fix_strs, e.g. fix 1 blah blah blah)
      self.fix_ids   = [ f.split()[1] for f in self.fix_strs ]
      # get unfix ids
      unfix_ids      = [ l.split()[1] for l in lines if "unfix" in l ]

      # if fix defined, remove minimize (i.e., assume it was from a step before)
      if len(self.fix_strs) > 1:
        for ii in range(len(self.fix_strs)):
          if self.fix_strs[ii].startswith("minimize"):
            del self.fix_strs[ii]
            del self.fix_ids[ii]
            break

      if False:
        print("FOR FILE:", self.log_file)
        print("Lines:", self.step_line_num, self.loop_line_num)
        print("ALL LINES:", lines)
        print("Fixes:", self.fix_strs)
        print("Unfix IDS:", unfix_ids)

      # if unfix active, remove the first instance of fix with that same fix id as the unfix
      for unfix_id in unfix_ids:
        for ii in range(len(self.fix_strs)):
          if unfix_id == self.fix_ids[ii]:
            del self.fix_ids[ii]
            del self.fix_strs[ii]
            break

      if False:
        print("Remainig Fixes:", self.fix_strs)
        print(self.df)

  # initializes the class
  def __init__(self, log_filei, step_line_numi, loop_line_numi):
    self.log_file = log_filei
    self.step_line_num = step_line_numi
    self.loop_line_num = loop_line_numi
    self.get_df()
    self.get_fixes()

# collect all log files
import glob
log_files = glob.glob("*/*.log")

# this list will hold all of the instances of Step class; one for each step in the protocol
steps     = []

# Next, we need to read in all of the data frames
for log_file in log_files:
  lines = ""
  with open(log_file) as f:
    lines = f.readlines()
  step_line_nums = [ n+1 for l, n in zip(lines, np.arange(len(lines))) if l.startswith("Step") ]
  loop_line_nums = [ n+1 for l, n in zip(lines, np.arange(len(lines))) if l.startswith("Loop") ]

  # remove all steps that are run=0
  idx_to_delete = []
  for i in range(len(step_line_nums)):
    if step_line_nums[i] == loop_line_nums[i]-2:
      idx_to_delete.append(i)
  for i in sorted(idx_to_delete, reverse=True):
    del step_line_nums[i]
    del loop_line_nums[i]

  for i in range(len(step_line_nums)):
    steps.append(Step(log_file, step_line_nums[i], loop_line_nums[i]))

fig, axs = plt.subplots(2, 2, figsize=(12,10), sharex=True, constrained_layout=True)

offset = 0
for step in steps:
  labeli = ';'.join(step.fix_strs)
  labeli = offset
  labeli = step.log_file
  axs[0, 0].plot([offset, offset+1], [step.df["P_all"].iat[0], step.df["P_all"].iat[-1]], label=labeli)
  axs[0, 1].plot([offset, offset+1], [step.df["P_all"].iat[0], step.df["P_all"].iat[-1]], label=labeli)
  axs[1, 0].plot([offset, offset+1], [step.df["P_pol"].iat[0], step.df["P_pol"].iat[-1]], label=labeli)
  axs[1, 1].plot([offset, offset+1], [step.df["P_fil"].iat[0], step.df["P_fil"].iat[-1]], label=labeli)
  offset += 1

axs[1, 0].set_xlabel("Step Number")
axs[1, 1].set_xlabel("Step Number")
axs[0, 1].set_ylabel("P_all")
axs[1, 0].set_ylabel("P_pol")
axs[1, 1].set_ylabel("P_fil")
axs[0, 0].legend(loc="best", fontsize=8, ncol=1)

fig.savefig("fig-pnorm.pdf")