move lib stuff to head so toksearch parallelization works; add script to dump info for experimental runs like textual summaries (since toksearch only works with shot data)

Author: +Joseph Abbate

lib/OMFITlib_fit.py OMFITlib_fit.py

@@ -2,7 +2,7 @@ data:
   shots: data/shots.npy
   tmin: 0
   tmax: 6000
-  time_step: 25
+  time_step: 20
   sql_sig_names: ['t_ip_flat','ip_flat_duration','topology','time_of_shot','run',
                   'poh','pech','pbeam','btor','btorsign','ip','ipsign', 'betanmax']
   scalar_sig_names: ['vloop','ip','ipsiptargt','bt',
@@ -29,28 +29,30 @@ data:
   efit_types: ['EFIT01','EFIT02','EFIT_CAKE01','EFIT_CAKE02']
   include_psirz: False
   include_rhovn: True
-  thomson_sig_names: [] #['density', 'temp']
+  thomson_sig_names: ['density', 'temp']
   include_thomson_uncertainty: True
   include_rt_thomson: True
-  cer_sig_names: [] #['temp','rot']
+  cer_sig_names: ['temp','rot']
   cer_type: 'CERQUICK'
   cer_realtime_channels: True
-  cer_rotation_units_of_khz: True
+  cer_rotation_units_of_krad: True
   zipfit_sig_names: ['trotfit','edensfit','etempfit','itempfit'] #,'idensfit']
 #['ftscrot','ftscpsin','ftsc1vld','ftsspsin','ftssrot','etscr','etscrin','etscrout','etsct','etsctin']
 # ['etste', 'etsne','etscr','etsct',
 #                'etstein', 'etsnein','etscrin', 'etsctin','etsinq', 'etsinprs',
 #                'etsteout', 'etsneout', 'etsqout', 'etsprsout',
 #                'etste', 'etsne','etscr','etsct']
   pcs_sig_names: []
-  include_radiation: False
+  include_radiation: True
   include_full_ech_data: True
   include_full_nb_data: True
   include_gas_valve_info: True
+  include_log_info: True
   num_x_points: 33
   trial_fits: [] #['spline_1d','linear_interp_1d','mtanh_1d', 'csaps_1d']
 logistics:
-  output_file: '/mnt/beegfs/users/abbatej/example.h5'
+  output_file: '/mnt/beegfs/users/abbatej/new_20ms_log_data/example.h5'
   overwrite_shots: True
   num_processes: 1
-  max_shots_per_run: 5000
+  max_shots_per_run: 1
+  print_errors: False

configs/example.yaml

@@ -5,8 +5,8 @@
 
 # for getting all plasma shots past a certain shot number
 if True:
-    minimum_shot = 165400
-    maximum_shot=174042
+    minimum_shot = 140000 #2010-ish
+    maximum_shot = 194528 #200000 #arbitrarily big
     query = """
         select shot
         from shots_type

dump_shots.py

@@ -31,7 +31,6 @@
 import os
 import time
 from scipy import interpolate, stats
-sys.path.append(os.path.join(os.path.dirname(__file__),'lib'))
 from transport_helpers import my_interp, standardize_time, Timer
 import fit_functions
 from plot_tools import plot_comparison_over_time, plot_2d_comparison
@@ -156,6 +155,7 @@
     print(f'Starting shot {shots[0]}-{shots[-1]}')
     sys.stdout.flush()
 
+    print('Writing summary SQL signals')
     # pipeline for SQL signals
     if len(cfg['data']['sql_sig_names'])>0:
         conn = connect_d3drdb()
@@ -185,6 +185,7 @@
                     else:
                         final_data[shot][sig_name]=record[sig]
 
+    print('Writing gas SQL signals')
     # pipeline for GAS
     if cfg['data']['include_gas_valve_info']:
         gas_sigs=['gas','valve']
@@ -203,14 +204,41 @@
         for record in records:
             for sig in gas_sigs:
                 shot=str(record['shot'])
-                tmp_dic[shot][sig].append(record[sig])
+                tmp_dic[shot][sig].append(str(record[sig]))
         with h5py.File(filename,'a') as final_data:
             for shot in tmp_dic:
                 final_data.require_group(shot)
                 for sig in gas_sigs:
                     sig_name=sig+'_sql'
                     final_data[shot][sig_name]=tmp_dic[shot][sig]
 
+    print('Writing log SQL signals')
+    # pipeline for LOGS
+    if cfg['data']['include_log_info']:
+        log_sigs=['text','topic','username']
+        conn = connect_d3drdb()
+        query="""SELECT shot,{}
+                 FROM entries
+                 WHERE shot in {}
+              """.format(
+            ','.join(log_sigs),
+            '({})'.format(','.join([str(elem) for elem in shots]))
+            )
+        pipeline = Pipeline.from_sql(conn, query)
+        records=pipeline.compute_serial()
+        tmp_dic={str(shot): {sig: [] for sig in log_sigs} for shot in shots}
+        for record in records:
+            for sig in log_sigs:
+                shot=str(record['shot'])
+                tmp_dic[shot][sig].append(str(record[sig]))
+        with h5py.File(filename,'a') as final_data:
+            for shot in tmp_dic:
+                final_data.require_group(shot)
+                for sig in log_sigs:
+                    sig_name=sig+'_sql'
+                    final_data[shot][sig_name]=tmp_dic[shot][sig]
+
+    print('Writing timebased signals')
     # pipeline for regular signals
     pipeline = Pipeline(shots)
 
@@ -461,21 +489,22 @@ def add_ech_info(record):
             record['ech_pwr_total']=standardize_time(record[f'ech_pwr_total_full']['data'],
                                                      record[f'ech_pwr_total_full']['times'],
                                                      record['standard_time'])
-            num_systems=record['ech_num_systems']['data']
-            record['ech_names']=[]
-            sigs_0d=['frequency','R','Z']
-            sigs_1d=['pwr','aziang','polang']
-            for key in sigs_0d+sigs_1d:
-                record[f'ech_{key}']=[]
-            for i in range(1,num_systems+1):
-                gyro=record[f'ech_name_{i}']['data'].upper()
-                record['ech_names'].append(gyro)
-                for key in sigs_0d:
-                    record[f'ech_{key}'].append(record[f'ech_{key}_{i}']['data'])
-                for key in sigs_1d:
-                    record[f'ech_{key}'].append(standardize_time(record[f'ech_{key}_{gyro}']['data'],
-                                                                record[f'ech_{key}_{gyro}']['times'],
-                                                                record['standard_time']))
+            if record['ech_num_systems'] is not None:
+                num_systems=record['ech_num_systems']['data']
+                record['ech_names']=[]
+                sigs_0d=['frequency','R','Z']
+                sigs_1d=['pwr','aziang','polang']
+                for key in sigs_0d+sigs_1d:
+                    record[f'ech_{key}']=[]
+                for i in range(1,num_systems+1):
+                    gyro=record[f'ech_name_{i}']['data'].upper()
+                    record['ech_names'].append(gyro)
+                    for key in sigs_0d:
+                        record[f'ech_{key}'].append(record[f'ech_{key}_{i}']['data'])
+                    for key in sigs_1d:
+                        record[f'ech_{key}'].append(standardize_time(record[f'ech_{key}_{gyro}']['data'],
+                                                                    record[f'ech_{key}_{gyro}']['times'],
+                                                                    record['standard_time']))
 
     if cfg['data']['include_full_nb_data']:
         @pipeline.map
@@ -724,7 +753,7 @@ def pcs_processing(record):
 
     with Timer():
         if cfg['logistics']['num_processes']>1:
-            records=pipeline.compute_spark(numparts=cfg['logistics']['num_processes'])
+            records=pipeline.compute_ray(numparts=cfg['logistics']['num_processes'])
         else:
             records=pipeline.compute_serial()
 
@@ -738,6 +767,9 @@ def pcs_processing(record):
                 if sig in final_data[shot]:
                     del final_data[shot][sig]
                 final_data[shot][sig]=record[sig]
-            for key in record['errors']:
-                print(key)
-                print(record['errors'][key]['traceback'].replace('\\n','\n'))
+                # print(sig)
+                # print(record[sig])
+            if cfg['logistics']['print_errors']:
+                for key in record['errors']:
+                    print(key)
+                    print(record['errors'][key]['traceback'].replace('\\n','\n'))

lib/fit_functions.py fit_functions.py

@@ -0,0 +1,63 @@
+"""
+module load omfit
+omfit run_dump.py
+
+(or copy/paste into an OMFIT command window and run from there)
+"""
+
+import h5py
+import numpy as np
+import time
+
+filename='/cscratch/abbatej/run_dump.h5'
+min_shot=140000
+max_shot=196512
+overwrite_runs=True
+
+# splitting up into batches prob not necessary in retrospect...
+# the thought was to not have too much stuff in RAM at a time
+max_shots_per_run=10000
+
+text_sigs=['text','topic','username']
+run_sigs=['brief']
+run_sig_string=','.join(run_sigs)
+text_string=','.join(text_sigs)
+
+shot_dividers=list(reversed(np.arange(min_shot,max_shot,max_shots_per_run)))
+
+prev_time=time.time()
+for i in range(len(shot_dividers)-1):
+    prev_time=time.time()
+    this_max_shot=shot_dividers[i]
+    this_min_shot=shot_dividers[i+1]
+    print(f"Starting {this_max_shot}-{this_min_shot}")
+    sig_info=OMFITrdb(f"""SELECT DISTINCT UPPER(run) as run
+    FROM shots
+    WHERE shot>={this_min_shot} AND shot<{this_max_shot}""",db='d3drdb',server='d3drdb',by_column=True)
+    runs=sig_info['run']
+    run_string='({})'.format(','.join(["'"+run+"'" for run in runs]))
+    tmp_dic={run: {sig: [] for sig in text_sigs} for run in runs}
+
+    sig_info=OMFITrdb(f"""SELECT UPPER(run) as run,{run_sig_string}
+    FROM runs
+    WHERE run in {run_string}""",db='d3drdb',server='d3drdb',by_column=True)
+    for i in range(len(sig_info['run'])):
+        for sig in run_sigs:
+            tmp_dic[str(sig_info['run'][i])][sig]=str(sig_info[sig][i])
+
+    sig_info=OMFITrdb(f"""SELECT UPPER(run) as run,{text_string}
+    FROM entries
+    WHERE run in {run_string} AND shot IS NULL""",db='d3drdb',server='d3drdb',by_column=True)
+    for i in range(len(sig_info['run'])):
+        for sig in text_sigs:
+            tmp_dic[str(sig_info['run'][i])][sig].append(str(sig_info[sig][i]))
+
+    with h5py.File(filename, 'a') as f:
+        for run in tmp_dic:
+            if overwrite_runs and run in f:
+                del f[run]
+            else:
+                f.require_group(run)
+                for sig in tmp_dic[run]:
+                    f[run][sig]=tmp_dic[run][sig]
+    print(f'Took {time.time()-prev_time:.2f}s')

lib/mtanh_fit.py mtanh_fit.py

@@ -25,7 +25,6 @@ def __enter__(self):
 
     def __exit__(self, *args):
         elapsed = time.time() - self.start
-        print('---------------------')
         print('---> Ran in {0:.2f} s'.format(elapsed))
 
 def standardize_time(old_signal,old_timebase,standard_times,