#177: changed output filenames, included beta* in import

kmod/gui2beta/get_kmod_files.py

@@ -1,66 +1,89 @@
-#!/afs/cern.ch/work/o/omc/anaconda/bin/python
+"""
+Module kmod.get_kmod_files
+-----------------------------
 
+Used by the BB-Gui to merge and copy kmod files from subfolders of the chosen folder into
+the current optics output folder.
+"""
 import os
 import sys
 from os.path import abspath, join, dirname, pardir
 new_path = abspath(join(dirname(abspath(__file__)), pardir, pardir))
 if new_path not in sys.path:
     sys.path.append(new_path)
 
-from Python_Classes4MAD import metaclass
 from optparse import OptionParser
-from tfs_files import tfs_file_writer
+from tfs_files import tfs_pandas
 import re
+from utils import logging_tools
 
+from gui2kmod import get_beta_star_filename, get_beta_filename
 
-def write_global_files(beam, kmod_dir, res_dir, mod_path):    
+LOG = logging_tools.get_logger(__name__)
+
+
+# Main Functions ###############################################################
+
+
+def merge_and_copy_kmod_output(beam, kmod_dir, res_dir, mod_path):
+    """ Merges the needed files into dataframes and saves them.
+
+    Args:
+        beam (str): Currently used beam.
+        kmod_dir (str): Parent folder of the kmod results
+        res_dir (str): Destination folder for the merged output
+        mod_path (str): Path to the model.
+    """
+    LOG.debug("Merging and copying of kmod data started.")
     pattern = re.compile(".*R[0-9]\." + beam.upper())
-    model_twiss = metaclass.twiss(mod_path)
-
-    ip_dir_names = []
-    for path,dirnames,files in os.walk(kmod_dir):
-        for dirname in dirnames:
-            if pattern.match(dirname):
-                ip_dir_names.append(dirname)
-
-    for plane in ("x", "y"):
-        uplane = plane.upper()
-        results_writer = tfs_file_writer.TfsFileWriter.open(join(res_dir, 'getkmodbeta' + plane + '.out'))
-        write_headers(results_writer, uplane)
-        bpm_names, bpm_s, betas, betas_std, betas_mdl = [], [], [], [], []
+    model_twiss = tfs_pandas.read_tfs(mod_path, index="NAME")
+
+    ip_dir_names = [d for _, dirs, _ in os.walk(kmod_dir) for d in dirs if pattern.match(d)]
+
+    # copy beta data
+    for plane in "xy":
+        new_data = tfs_pandas.TfsDataFrame()
         for ip_dir_name in ip_dir_names:
-            path = os.path.join(kmod_dir, ip_dir_name, 'getkmodbeta' + plane + '.out')
-            try:
-                data = metaclass.twiss(path)
-            except IOError:
-                print "Cannot find kmod data in " + path + ", won't copy those files."
-                continue
-            print("Kmod data found in: {}".format(path))
-            for bpm_name in data.NAME:
-                if bpm_name in model_twiss.NAME:
-                    index = data.indx[bpm_name]
-                    bpm_names.append(bpm_name)
-                    betas.append(getattr(data, "BET" + uplane)[index])
-                    betas_std.append(getattr(data, "BET" + uplane + "STD")[index])
-                    model_index = model_twiss.indx[bpm_name]
-                    betas_mdl.append(getattr(model_twiss, "BET" + uplane)[model_index])
-                    bpm_s.append(getattr(model_twiss, "S")[model_index])
-
-        for i in range(len(bpm_names)):
-            results_writer.add_table_row([bpm_names[i], bpm_s[i], 0, betas[i],
-                                          betas_std[i], betas_mdl[i], 0, 0, 0 ])
-        results_writer.write_to_file()
-
-
-def write_headers(resuts_writers, uplane):
-    resuts_writers.set_column_width(20)
-    resuts_writers.add_column_names(
-        ['NAME', 'S', 'COUNT', 'BET' + uplane, 'STDBET' + uplane,
-         'BET' + uplane + 'MDL', 'ERRBET' + uplane, 'BET' + uplane + 'RES',
-         'BET' + uplane + 'STDRES']
-    )
-    resuts_writers.add_column_datatypes(['%s', '%le', '%le', '%le', '%le',
-                                         '%le', '%le', '%le', '%le'])
+            src = join(kmod_dir, ip_dir_name, get_beta_filename(plane))
+            data = _load_source_data(src, "NAME")
+            if data is not None:
+                new_data = new_data.append(data.loc[data.index.isin(model_twiss.index), :])
+        new_data["S"] = model_twiss.loc[new_data.index, "S"]
+        dst = join(res_dir, get_beta_merged_filename(plane))
+        tfs_pandas.write_tfs(dst, new_data, save_index="NAME")
+
+    # copy beta* data
+    new_data = tfs_pandas.TfsDataFrame()
+    for ip_dir_name in ip_dir_names:
+        src = join(kmod_dir, ip_dir_name, get_beta_star_filename())
+        data = _load_source_data(src)
+        new_data = new_data.append(data)
+    dst = join(res_dir, get_beta_star_merged_filename())
+    tfs_pandas.write_tfs(dst, new_data)
+
+
+def get_beta_star_merged_filename():
+    """ Outputfilename of the merged betastar file. """
+    return "getkmodbetastar.out"
+
+
+def get_beta_merged_filename(plane):
+    """ Outputfilename of the merged beta file. """
+    return get_beta_filename(plane)
+
+
+# Private Functions ############################################################
+
+
+def _load_source_data(src, index=None):
+    data = None
+    try:
+        data = tfs_pandas.read_tfs(src, index=index)
+    except IOError:
+        LOG.warn("Cannot find kmod data in '{:s}', won't copy those files.".format(src))
+    else:
+        LOG.warn("Loaded kmod data from '{:s}'.".format(src))
+    return data
 
 
 def parse_args():
@@ -74,19 +97,17 @@ def parse_args():
                             action='store', type='string', dest='res_dir')
     parser.add_option('-m', '--model_path',
                             help='Specify path to current model',
-                            action='store', type='string', dest='mod')
+                            action='store', type='string', dest='mod_path')
     parser.add_option('-b', '--beam',
                             help='define beam used: b1 or b2',
                             action='store', type='string', dest='beam')
     (options, _) = parser.parse_args()
     return options
 
 
-if __name__ == '__main__':
-    options = parse_args()
-    beam = options.beam
-    kmod_dir = options.kmod_dir
-    res_dir  = options.res_dir
-    mod_path = options.mod
+# Script Mode ##################################################################
 
-    write_global_files(beam, kmod_dir, res_dir, mod_path)
+
+if __name__ == '__main__':
+    opt = parse_args()
+    merge_and_copy_kmod_output(opt.beam, opt.kmod_dir, opt.res_dir, opt.mod_path)

kmod/gui2beta/gui2kmod.py

@@ -22,10 +22,10 @@
 
 CURRENT_PATH = os.path.abspath(os.path.dirname(__file__))
 
-
-# TODO: (Long term) Think about the accelerator class here for positions and Ks
-# TODO: Short term: Use a logger for logging
-# TODO: Short term: Use tfs_pandas instead of metaclass
+# TODO: Short term: Think about the accelerator class here for positions and Ks
+# TODO: Immediately: Use a logger for logging
+# TODO: Immediately: get rid of repetive code and use loops and functions
+# TODO: Immediately: Use tfs_pandas instead of metaclass
 
 def parse_args():
     parser = argparse.ArgumentParser()
@@ -304,7 +304,7 @@ def run_analysis_simplex(path, beam, magnet1, magnet2, hor_bstar, vert_bstar, wa
                                       fity_def, erry_def, ek, ek, cminus, vert_bstar, waist,
                                       (magnet1 + '-' + magnet2 + '.' + beam) + '.Y', log, logfile)
 
-    results = tfs_file_writer.TfsFileWriter.open(os.path.join(path, '%s.results' % (magnet1 + magnet2 + beam)))
+    results = tfs_file_writer.TfsFileWriter.open(os.path.join(path, get_results_filename()))
     results.set_column_width(15)
     results.add_column_names(
         ['LABEL', 'BETAWAIST', 'BETAWAIST_ERR', 'WAIST', 'WAIST_ERR', 'BETA_AV_FOC', 'BETA_AV_FOC_ERR', 'BETA_AV_DEF',
@@ -337,7 +337,7 @@ def calc_beta_instr(path, magnet1, magnet2, beam, instr, log, logfile, twiss):
 
         L_star_position = Magnet_definitions.LstarPosition(magnet1, magnet2, beam, twiss)
 
-        result_waist = metaclass.twiss(os.path.join(path, '%s.results' % (magnet1 + magnet2 + beam)))
+        result_waist = metaclass.twiss(os.path.join(path, get_results_filename()))
         beta_waist = result_waist.BETAWAIST
         beta_waist_err = result_waist.BETAWAIST_ERR
         waist = result_waist.WAIST
@@ -387,15 +387,15 @@ def calc_beta_instr(path, magnet1, magnet2, beam, instr, log, logfile, twiss):
         err_y = (abs(np.nanmax(beta_err_y, axis=0) - np.nanmin(beta_err_y, axis=0))) / 2.
 
         if name == 'BPM':
-            xdata = tfs_file_writer.TfsFileWriter.open(os.path.join(path, 'getkmodbetax.out'))
+            xdata = tfs_file_writer.TfsFileWriter.open(os.path.join(path, get_beta_filename("x")))
         else:
             xdata = tfs_file_writer.TfsFileWriter.open(os.path.join(path, 'Beta_%s_X.out' % name))
         xdata.set_column_width(20)
         xdata.add_column_names(['NAME', 'S', 'COUNT', 'BETX', 'BETXSTD', 'BETXMDL', 'MUXMDL', 'BETXRES', 'BETXSTDRES'])
         xdata.add_column_datatypes(['%s', '%le', '%le', '%le', '%le', '%le', '%le', '%le', '%le'])
 
         if name == 'BPM':
-            ydata = tfs_file_writer.TfsFileWriter.open(os.path.join(path, 'getkmodbetay.out'))
+            ydata = tfs_file_writer.TfsFileWriter.open(os.path.join(path, get_beta_filename("y")))
         else:
             ydata = tfs_file_writer.TfsFileWriter.open(os.path.join(path, 'Beta_%s_Y.out' % name))
         ydata.set_column_width(20)
@@ -417,13 +417,13 @@ def calc_beta_star(path, magnet1, magnet2, beam, lstar, twiss):
     if Magnet_definitions.FindParentBetweenMagnets(magnet1, magnet2, 'OMK', beam, twiss):
 
         results_write = tfs_file_writer.TfsFileWriter.open(
-            os.path.join(path, '%s.beta_star.results' % (magnet1 + magnet2 + beam)))
+            os.path.join(path, get_beta_star_filename()))
         results_write.set_column_width(20)
         results_write.add_column_names(
             ['LABEL', 'BETASTAR', 'BETASTAR_ERR', 'WAIST', 'WAIST_ERR', 'BETAWAIST', 'BETAWAIST_ERR', 'PHASEADV', 'ERRPHASEADV'])
         results_write.add_column_datatypes(['%s', '%le', '%le', '%le', '%le', '%le', '%le', '%le', '%le'])
 
-        results = metaclass.twiss(os.path.join(path, '%s.results' % (magnet1 + magnet2 + beam)))
+        results = metaclass.twiss(os.path.join(path, get_results_filename()))
         beta_w = results.BETAWAIST
 
         for i, b_w in enumerate(beta_w):
@@ -618,6 +618,19 @@ def _main():
     logdata.close()
 
 
+def get_beta_filename(plane):
+    return "getkmodbeta{:s}.out".format(plane)
+
+
+def get_beta_star_filename():
+    return "beta_star.out"
+
+
+def get_results_filename():
+    return "results.out"
+
+
+
 if __name__ == '__main__':
     _main()