Skip to content

Commit

Permalink
Add types to all CLI arguments
Browse files Browse the repository at this point in the history
  • Loading branch information
@bcaddy
bcaddy committed Aug 14, 2024
1 parent 68b3ad9 commit 05d1a6c
Show file tree
Hide file tree
Showing 2 changed files with 77 additions and 151 deletions.
16 changes: 8 additions & 8 deletions iseult.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,35 +8,35 @@
parser = argparse.ArgumentParser(description='Plotting program for Tristan-MP files.')

parser.add_argument('-n', nargs = '?',# dest='accumulate', action='store_const',
const=-1, default=-1,
const=-1, default=-1, type=int,
help='Maximum file # to consider')
parser.add_argument('-framerate', nargs = '?',# dest='accumulate', action='store_const',
const=10, default=10,
const=10, default=10, type=int,
help='FPS for the movie')

parser.add_argument('-outmovie', nargs = '?',# dest='accumulate', action='store_const',
const='out.mov', default='out.mov',
const='out.mov', default='out.mov', type=str,
help='FPS for the movie')

parser.add_argument('-O', nargs = '+',# dest='accumulate', action='store_const',
default=[''],
default=[''], type=str,
help='Directory Iseult will open. Default is output')

parser.add_argument('-p', nargs = '?',# dest='accumulate', action='store_const',
const='Default', default='Default',
const='Default', default='Default', type=str,
help='''Open Iseult with the given saved view.
If the name of view contains whitespace,
either it must be enclosed in quotation marks or given
with whitespace removed. Name is case sensitive.''')
parser.add_argument("-b", help="Run Iseult from bash script. Makes a movie.",
parser.add_argument("-b", help="Run Iseult from bash script. Makes a movie.", type=bool,
action="store_true")

parser.add_argument("-name", nargs = '+',# dest='accumulate', action='store_const',
default=[''],
default=[''], type=str,
help='Plot Title')

parser.add_argument("--wait", help="Wait until current simulation is finished before making movie.",
action="store_true")
action="store_true", type=bool)

parser.add_argument("-e",
"--electron-spectra",
Expand Down
212 changes: 69 additions & 143 deletions src/main_app.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -432,7 +432,7 @@ def __init__(self, parent, param, canvas = None):
self.param.attach(self)

def OnReload(self, *args):
self.parent.ReloadPath()
_ = self.parent.checkAndFindFilePaths()
self.parent.RenewCanvas()

def OnRefresh(self, *args):
Expand Down Expand Up @@ -2115,155 +2115,81 @@ def SaveIseultState(self, cfgfile, cfgname):
with open(cfgfile, 'w') as cfgFile:
cfgFile.write(yaml.safe_dump(cfgDict))

def ReloadPath(self):
""" This function updates the current pathdictionary"""
dirlist = os.listdir(self.dirname)

if int(self.cmd_args.n)!=-1:
self.CheckMaxNPopUp()

# Create a dictionary of all the paths to the files
self.PathDict = {'Flds': [], 'Prtl': [], 'Param': [], 'Spect': []}

# create a bunch of regular expressions used to search for files
f_re = re.compile('flds.tot.*')
prtl_re = re.compile('prtl.tot.*')
s_re = re.compile('spect.*')
param_re = re.compile('param.*')
self.PathDict['Flds']= list(filter(f_re.match, os.listdir(self.dirname)))
self.PathDict['Flds'].sort()
self.PathDict['Prtl']= list(filter(prtl_re.match, os.listdir(self.dirname)))
self.PathDict['Prtl'].sort()
self.PathDict['Spect']= list(filter(s_re.match, os.listdir(self.dirname)))
self.PathDict['Spect'].sort()
self.PathDict['Param']= list(filter(param_re.match, os.listdir(self.dirname)))
self.PathDict['Param'].sort()

### iterate through the Paths and just get the .nnn number
if len(self.PathDict['Param']) > 0:
self.length_of_outfiles = len(self.PathDict['Param'][0].split('.')[-1])
else:
self.length_of_outfiles = 3

for key in self.PathDict.keys():
for i in range(len(self.PathDict[key])):
try:
self.PathDict[key][i] = int(self.PathDict[key][i].split('.')[-1])
except ValueError:
self.PathDict[key].pop(i)
except IndexError:
pass

### GET THE NUMBERS THAT HAVE ALL 4 FILES:

allFour = set(self.PathDict['Param'])
for key in self.PathDict.keys():
allFour &= set(self.PathDict[key])
allFour = sorted(allFour)

if int(self.cmd_args.n) != -1 and len(allFour)>0:
while allFour[-1] > int(self.cmd_args.n) and len(allFour)>0:
allFour.pop(-1)

# Rebuild the pathdict only with files that have all 4 things
self.PathDict['Flds'] = ['flds.tot.'+str(elm).zfill(self.length_of_outfiles) for elm in allFour]
self.PathDict['Prtl'] = ['prtl.tot.'+str(elm).zfill(self.length_of_outfiles) for elm in allFour]
self.PathDict['Spect'] = ['spect.'+str(elm).zfill(self.length_of_outfiles) for elm in allFour]
self.PathDict['Param'] = ['param.'+str(elm).zfill(self.length_of_outfiles) for elm in allFour]
self.TimeStep.setMax(len(self.PathDict['Flds']))
self.playbackbar.slider.config(to =(len(self.PathDict['Flds'])))
if self.MainParamDict['Reload2End']:
self.TimeStep.value = len(self.PathDict['Flds'])
self.playbackbar.slider.set(self.TimeStep.value)
self.shock_finder() #hack

def CheckMaxNPopUp(self):
MaxNDialog(self)
def pathOK(self):
""" Test to see if the current path contains tristan files
using regular expressions, then generate the lists of files
to iterate over"""
dirlist = os.listdir(self.dirname)
if 'output' in dirlist:
self.dirname = os.path.join(self.dirname, 'output')

is_okay = True

# Create a dictionary of all the paths to the files
self.PathDict = {'Flds': [], 'Prtl': [], 'Param': [], 'Spect': []}

# create a bunch of regular expressions used to search for files
f_re = re.compile(r'flds.tot.([\d\w\.]+)(?<!xdmf)$') # !(*.xdmf)
prtl_re = re.compile(r'prtl.tot.*')
s_re = re.compile(r'spect.*')
param_re = re.compile(r'param.*')

self.PathDict['Flds']= list(filter(f_re.match, os.listdir(self.dirname)))
if len(self.PathDict['Flds']) == 0 and pathlib.Path(self.dirname+'/flds').is_dir():
self.PathDict['Flds']= list(filter(f_re.match, os.listdir(self.dirname+'/flds')))
self.PathDict['Flds'].sort()

self.PathDict['Prtl']= list(filter(prtl_re.match, os.listdir(self.dirname)))
if len(self.PathDict['Prtl']) == 0 and pathlib.Path(self.dirname+'/prtl').is_dir():
self.PathDict['Prtl']= list(filter(prtl_re.match, os.listdir(self.dirname+'/prtl')))
self.PathDict['Prtl'].sort()

self.PathDict['Spect']= list(filter(s_re.match, os.listdir(self.dirname)))
if len(self.PathDict['Spect']) == 0:
s_re = re.compile(r'spec.*')
self.PathDict['Spect']= list(filter(s_re.match, os.listdir(self.dirname)))
if len(self.PathDict['Spect']) >= 1 and pathlib.Path(self.dirname+'/spec').is_dir():
self.PathDict['Spect']= list(filter(s_re.match, os.listdir(self.dirname+'/spec')))
self.PathDict['Spect'].sort()

self.PathDict['Param']= list(filter(param_re.match, os.listdir(self.dirname)))
self.PathDict['Param'].sort()

### iterate through the Paths and just get the .nnn number
if len(self.PathDict['Param']) > 0:
self.length_of_outfiles = len(self.PathDict['Param'][0].split('.')[-1])
def __findFilesInSubdir(self, file_glob: str, subdir: str) -> list[pathlib.Path]:
file_list = list(self.dirname.glob(file_glob))
if len(file_list) == 0:
file_list = list((self.dirname/subdir).glob(file_glob))

return file_list

def checkAndFindFilePaths(self, reload_mode: bool = False) -> bool:
# Convert to pathlib.Path
self.dirname = pathlib.Path(self.dirname)

# Check for "output" directory and update path if it's found
if (self.dirname/'output').is_dir():
self.dirname /= 'output'

# Check if data is from Tristan v1 or v2 and set file names accordingly
if len(list(self.dirname.glob("param.*"))) > 0:
# This is Tristan v1 data
param_name = 'param'
spectra_name = 'spect'
elif len(list(self.dirname.glob("params.*"))) > 0:
# This is Tristan v2 data
param_name = 'params'
spectra_name = 'spec'
else:
self.length_of_outfiles = 3
# This directory does not contain any Tristan files
return False

for key in self.PathDict.keys():
for i in range(len(self.PathDict[key])):
try:
self.PathDict[key][i] = int(self.PathDict[key][i].split('.')[-1])
except ValueError:
self.PathDict[key].pop(i)
except IndexError:
pass
# Load each list of files, searching in subdirectories if needed
param_files = list(self.dirname.glob(f"{param_name}.*"))
flds_files = self.__findFilesInSubdir('flds.tot.*', 'flds')
prtl_files = self.__findFilesInSubdir('prtl.tot.*', 'flds')
spectra_files = self.__findFilesInSubdir(f'{spectra_name}.*', 'flds')

### GET THE NUMBERS THAT HAVE ALL 4 FILES:
# Strip out .xdmf files
flds_files = [path for path in flds_files if not path.suffix == '.xdmf']

allFour = set(self.PathDict['Param'])
for key in self.PathDict.keys():
allFour &= set(self.PathDict[key])
allFour = sorted(allFour)
# Find which files are in all four lists
intersection = set.intersection(set([filename.suffix for filename in param_files]),
set([filename.suffix for filename in flds_files]),
set([filename.suffix for filename in prtl_files]),
set([filename.suffix for filename in spectra_files]))

# Check that there is at least one complete set of data. If not then return early
if len(intersection) == 0:
return False

if int(self.cmd_args.n) != -1 and len(allFour)>0:
while allFour[-1] > int(self.cmd_args.n) and len(allFour)>0:
allFour.pop(-1)
# Limit us to the first N files if the relevant CLI flag is set
if self.cmd_args.n != -1:
intersection = sorted(intersection)[:self.cmd_args.n]

is_okay = len(allFour)>0
# Rebuild the pathdict only with files that have all 4 things
# Reduce file path lists to just the complete datasets, sort, and assign to member variables
self.PathDict['Param'] = sorted([path for path in param_files if path.suffix in intersection])
self.PathDict['Flds'] = sorted([path for path in flds_files if path.suffix in intersection])
self.PathDict['Prtl'] = sorted([path for path in prtl_files if path.suffix in intersection])
self.PathDict['Spect'] = sorted([path for path in spectra_files if path.suffix in intersection])

self.PathDict['Flds'] = ['flds.tot.'+str(elm).zfill(self.length_of_outfiles) for elm in allFour]
self.PathDict['Prtl'] = ['prtl.tot.'+str(elm).zfill(self.length_of_outfiles) for elm in allFour]
self.PathDict['Spect'] = ['spect.'+str(elm).zfill(self.length_of_outfiles) for elm in allFour]
self.PathDict['Param'] = ['param.'+str(elm).zfill(self.length_of_outfiles) for elm in allFour]
if is_okay:
if reload_mode:
if int(self.cmd_args.n)!=-1:
self.CheckMaxNPopUp()
else:
self.NewDirectory = True
self.TimeStep.setMax(len(self.PathDict['Flds']))
self.playbackbar.slider.config(to =(len(self.PathDict['Flds'])))
if self.MainParamDict['Reload2End']:
self.TimeStep.value = len(self.PathDict['Flds'])
self.playbackbar.slider.set(self.TimeStep.value)
self.shock_finder() #hack
self.movie_dir = ''

return is_okay
self.TimeStep.setMax(len(self.PathDict['Flds']))
self.playbackbar.slider.config(to =(len(self.PathDict['Flds'])))
if self.MainParamDict['Reload2End']:
self.TimeStep.value = len(self.PathDict['Flds'])
self.playbackbar.slider.set(self.TimeStep.value)
self.shock_finder()

return True

def OnOpen(self, e = None):
"""open a file"""
Expand All @@ -2277,7 +2203,7 @@ def OnOpen(self, e = None):

else:
self.dirname = tmpdir
if not self.pathOK():
if not self.checkAndFindFilePaths():
# p = MyDalog(self, 'Directory must contain either the output directory or all of the following: \n flds.tot.*, ptrl.tot.*, params.*, spect.*', title = 'Cannot find output files')
# self.wait_window(p.top)
self.findDir()
Expand All @@ -2288,7 +2214,7 @@ def ResetSession(self, e = None):
"""open a file"""
if int(self.cmd_args.n) != -1:
self.CheckMaxNPopUp()
self.pathOK()
self.checkAndFindFilePaths()
self.ReDrawCanvas()

def findDir(self, dlgstr = 'Choose the directory of the output files.'):
Expand All @@ -2301,11 +2227,11 @@ def findDir(self, dlgstr = 'Choose the directory of the output files.'):
self.dir_opt['mustexist'] = True
self.dir_opt['parent'] = self

if not self.pathOK():
if not self.checkAndFindFilePaths():
tmpdir = filedialog.askdirectory(title = dlgstr, **self.dir_opt)
if tmpdir != '':
self.dirname = tmpdir
if not self.pathOK():
if not self.checkAndFindFilePaths():
# p = MyDialog(self, 'Directory must contain either the output directory or all of the following: \n flds.tot.*, ptrl.tot.*, params.*, spect.*', title = 'Cannot find output files')
# self.wait_window(p.top)
self.findDir()
Expand Down Expand Up @@ -3133,8 +3059,8 @@ def ReDrawCanvas(self, keep_view = True):
self.SubPlotList[i][j].DrawGraph()

if self.MainParamDict['ShowTitle']:
tmpstr = self.PathDict['Prtl'][self.TimeStep.value-1].split('.')[-1]
self.f.suptitle(os.path.abspath(self.dirname)+ '/*.'+tmpstr+r' at time t = %d $\omega_{pe}^{-1}$' % round(self.DataDict['time']), size = 15)
tmpstr = self.PathDict['Prtl'][self.TimeStep.value-1].suffix
self.f.suptitle(os.path.abspath(self.dirname)+ '/*'+tmpstr+r' at time t = %d $\omega_{pe}^{-1}$' % round(self.DataDict['time']), size = 15)
if keep_view:
self.LoadView()

Expand Down Expand Up @@ -3234,7 +3160,7 @@ def RefreshCanvas(self, keep_view = True):
self.SubPlotList[i][j].RefreshGraph()

if self.MainParamDict['ShowTitle']:
tmpstr = self.PathDict['Prtl'][self.TimeStep.value-1].split('.')[-1]
tmpstr = self.PathDict['Prtl'][self.TimeStep.value-1].suffix
self.f.suptitle(os.path.abspath(self.dirname)+ '/*.'+tmpstr+r' at time t = %d $\omega_{pe}^{-1}$' % round(self.DataDict['time']), size = 15)

if keep_view:
Expand Down

0 comments on commit 05d1a6c

Please sign in to comment.