1D_plot_results.py

#!/usr/bin/env python3
import argparse as ap
import os

import numpy as np
import matplotlib.pyplot as plt

from matplotlib import cm
import matplotlib.colors as colors

def main():
  # Setup command line arguements.
  parser = ap.ArgumentParser(description='Plots the scalar flux generated by DOMA for 2D transport problems.')
  parser.add_argument('-i', type=str, dest='file_name', required=True, help='The path of the input file (assuming the outputs are in the same directory).')
  parser.add_argument('-l', '--log-scale', action='store_true', dest='log_scale', help='Whether the plot is log scale or not..')

  cli_args = parser.parse_args()
  dir_path = os.path.dirname(os.path.realpath(cli_args.file_name))
  input_name = os.path.splitext(os.path.basename(cli_args.file_name))[0]

  dim_file = open(dir_path + "/" + input_name + "_dims.txt", "r")
  dim_x = int(str(dim_file.readline()).replace("num_x: ", "").replace("\n", ""))
  grps = int(str(dim_file.readline()).replace("num_g: ", "").replace("\n", ""))

  x_vals = np.loadtxt(dir_path + "/" + input_name + "_meshx.txt")

  for grp in range(grps):
    raw_flux = np.loadtxt(dir_path + "/" + input_name + "_g" + str(grp) + "_flux.txt")
    if cli_args.log_scale == True:
      plt.plot(x_vals, raw_flux)
      plt.yscale('log')
      plt.ylabel('Group ' + str(grp) + ' Scalar Flux (s$^{-1}$)')
      plt.xlabel('x (cm)')
    else:
      plt.plot(x_vals, raw_flux)
      plt.ylabel('Flux (s$^{-1}$)')
      plt.xlabel('x (cm)')

    plt.savefig(dir_path + "/" + input_name + "_g" + str(grp) + "_flux.png", format='png')
    plt.show()
    plt.close()

if __name__ == "__main__":
    main()