From f3a247e2fb693470f973939cb0df5239bc98ef0a Mon Sep 17 00:00:00 2001
From: Jeff Candy <candy@jaamweb.com>
Date: Wed, 4 Dec 2024 21:05:18 -0800
Subject: [PATCH] Updates to cgyro scripts

---
 cgyro/bin/cgyro_json  |  4 ++--
 cgyro/bin/cgyrodb     | 47 ++++++++++++++++++++++++++++++-------------
 f2py/requirements.txt |  1 +
 3 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/cgyro/bin/cgyro_json b/cgyro/bin/cgyro_json
index 68bd7a536..dd6298a47 100755
--- a/cgyro/bin/cgyro_json
+++ b/cgyro/bin/cgyro_json
@@ -25,14 +25,14 @@ def opts():
         description="CGYRO localdump utility",
         epilog=textwrap.dedent(mytext))
     
-    parser.add_argument('-dir',
+    parser.add_argument('-e',
                         help="working directory (default: cwd)",
                         type=str,
                         default='.')
     
     args=parser.parse_args()
 
-    return args.dir
+    return args.e
 
 mydir = opts()
 
diff --git a/cgyro/bin/cgyrodb b/cgyro/bin/cgyrodb
index 38751beb1..dacd501c3 100755
--- a/cgyro/bin/cgyrodb
+++ b/cgyro/bin/cgyrodb
@@ -13,6 +13,8 @@ from cgyro.data import cgyrodata
 
 pd.set_option('display.max_rows', None)
 
+ignore = 'eslshared'
+
 # Command line option parser
 def opts():
 
@@ -27,20 +29,32 @@ def opts():
         description="CGYRO database utility",
         epilog=textwrap.dedent(mytext))
 
-    parser.add_argument('-mode',
-                        help='Mode switch (data,update)',
-                        type=str,
-                        default='data')
+    parser.add_argument('-ref',
+                        help='Refresh directories',
+                        action='store_true')
+    parser.add_argument('-json',
+                        help='Generate IMAS and localdump files',
+                        action='store_true')
+    parser.add_argument('-db',
+                        help='Load the DB and print',
+                        action='store_true')
+    parser.add_argument('-flux',
+                        help='Add flux to DB',
+                        action='store_true')
 
     args=parser.parse_args()
 
-    return args.mode
+    return args.ref,args.json,args.db,args.flux
 
-mode = opts()
+ref,json,db,doflux = opts()
 
 meta = {}
 
-def gendict(sim,mode):
+if (not ref) and (not json) and (not db):
+    print('Type "cgyrodb -h" for help')
+    sys.exit()
+    
+def gendict(sim,doflux):
 
    mydict = {}
 
@@ -60,7 +74,7 @@ def gendict(sim,mode):
       mydict['a/Ln1'] = 0.0
       mydict['a/Lt1'] = 0.0
 
-   if mode == 'flux':
+   if doflux:
       sim.getflux()
       y = np.sum(sim.ky_flux,axis=(2,3))   
       for ispec in range(sim.n_species):
@@ -73,8 +87,6 @@ def gendict(sim,mode):
 
    return mydict
 
-ignore = 'eslshared'
-
 # First locate all directories
 y = []
 for root,xd,xf in os.walk('./'):
@@ -82,16 +94,23 @@ for root,xd,xf in os.walk('./'):
       if ignore not in root:
           y.append(root)
 
-if mode == 'update':
+if ref:
    # Run through directories and update
    for mdir in y:
-      print('Updating '+mdir)
+      print('Refreshing '+mdir)
       os.system('cgyro -t '+mdir)
-else:
+
+if json:
+   # Run through directories and update
+   for mdir in y:
+      print('Dumping '+mdir)
+      os.system('cgyro_json -e '+mdir)
+
+if db:
    # Create master dictionary "meta"
    for mdir in tqdm(y):
       sim = cgyrodata(mdir+'/',silent=True)
-      meta[mdir]=gendict(sim,mode)
+      meta[mdir]=gendict(sim,doflux)
 
    df = pd.DataFrame(meta)
    dfpart = df.T[['n_n','n_radial','n_theta','q','s','kappa','a/Ln0','a/Lt0','a/Ln1','a/Lt1','nuee']]
diff --git a/f2py/requirements.txt b/f2py/requirements.txt
index 7e99da002..dc04f0530 100644
--- a/f2py/requirements.txt
+++ b/f2py/requirements.txt
@@ -5,3 +5,4 @@ matplotlib
 fortranformat
 scikit-image
 build
+pandas