process_output_rmf.py

#! /usr/bin/env python


from __future__ import print_function
import argparse

import IMP
import IMP.rmf
import RMF
import difflib


parser = argparse.ArgumentParser(
    description='Process output data file saved as rmfs. It has two modality: '
    'print selected fields for all lines or print a particular line '
    'where a filed has a given value.')
parser.add_argument(
    '-f',
    action="store",
    dest="filename",
    help="file name to process")
parser.add_argument(
    '-s',
    dest="fields",
    nargs="+",
    help="Specify all fields to be printed. Multiple flags will append a "
         "list of fields to be printed")
parser.add_argument(
    '-p',
    action="store_true",
    dest="print_fields",
    default=False,
    help="print the fields contained in the file")
parser.add_argument(
    '-n',
    action="store",
    dest="nframe",
    default=None,
    help="Print frame number n")

result = parser.parse_args()


# open the file
if result.filename is not None:
    try:
        # let's see if that is an rmf file
        rh = RMF.open_rmf_file_read_only(result.filename)
        cat = rh.get_category('stat')
        rmf_klist = rh.get_keys(cat)
        rmf_names_keys = dict([(rh.get_name(k), k) for k in rmf_klist])
        klist = rmf_names_keys.keys()
        del rh
    except IOError:
        raise IOError("Not an RMF file")
else:
    raise ValueError("No file name provided. Use -h for help")


# print the keys
if result.print_fields:
    for key in klist:
        print(key)


# the field string matching is by default strict, i.e., the input string
# must be the same as the one in the file
match_strictness = 1.0

# print the queried fields
if result.fields is not None:
    field_list = []
    # check whether the fields exist and convert them to best matching existing
    # field names
    for field in result.fields:
        found_entries = difflib.get_close_matches(
            field,
            klist,
            1,
            match_strictness)
        if len(found_entries) == 0:
            raise ValueError("field " + field + " non found")
        else:
            field_list.append(found_entries[0])

    # print comment line
    s0 = ' '.join(["%20s" % (field) for field in field_list])
    print("# " + s0)

    # print fields values
    rh = RMF.open_rmf_file_read_only(result.filename)

    for frame_number in range(rh.get_number_of_frames()):
        IMP.rmf.load_frame(rh, RMF.FrameID(frame_number))
        rn = rh.get_root_node()
        for name in klist:
            s0 = ' '.join("%20s" % (str(rn.get_value(rmf_names_keys[name])))
                          for name in field_list)

        if not result.nframe:
            print("> " + s0)
        else:
            print(str(frame_number) + " > " + s0)
    del rh

# print given frame number
if result.nframe is not None:
    # print comment line
    # print fields values
    rh = RMF.open_rmf_file_read_only(result.filename)
    IMP.rmf.load_frame(rh, RMF.FrameID(int(result.nframe)))
    for k in klist:
        s0 = str(k)+" "+str(rh.get_root_node().get_value(rmf_names_keys[k]))
        print(s0)
    del rh