Hide
Siafoo – the intersection of pastebin, help desk, version control, and social networking Join Siafoo Now or Learn More

Stou's I/O and sorting assignment

Revision 1 vs. Revision 2

Legend:

Unmodified
Added
Removed
  • Code

    r1 r2  
    1919 
    20205. (bonus) Implement your own sorting method and compare to STL method. 
    2121 
    22226. (bonus) Try some of the Boost libraries 
    2323 
     247. (bonus) Accept command line arguments to pick algorithm, sort order, data 
     25   points. Try Boost Program Options 
     26 
    2427Python suggestions: 
    2528 
    26291. Use the standard sorting routine and define your own comparator function 
    2730 
    2831Oh and I suggest sticking to the standard libraries (STL and Python)... if you 
    2932decide to use stuff from Boost and numpy it would probably be good to implement 
    3033your own too and compare them. 
    3134 
    3235""" 
    3336import sys 
     37import numpy as np 
    3438 
    35 def dist(x, y, z, x0, y0, z0): 
    36     return ((x - x0)**2 + (y - y0)**2 + (z - z0)**2)**0.5 
     39def set_options(opt_list, func_dict): 
     40    opt_dict = {} 
     41 
     42    if '-xyz' in opt_list: 
     43        i = opt_list.index('-xyz') 
     44        xyz = float(opt_list[i+1]), float(opt_list[i+2]), float(opt_list[i+3]) 
     45        opt_dict['xyz0'] = xyz 
     46    else: 
     47        opt_dict['xyz0'] = None 
     48 
     49    if '-d' in opt_list: 
     50        i = opt_list.index('-d') 
     51        opt_dict['distfunc'] = func_dict[opt_list[i+1]] 
     52    else: 
     53        opt_dict['distfunc'] = calc_dist1 
     54 
     55    if '-s' in opt_list: 
     56        i = opt_list.index('-s') 
     57        opt_dict['sortfunc'] = func_dict[opt_list[i+1]] 
     58    else: 
     59        opt_dict['sortfunc'] = dist_sort1 
     60 
     61    return opt_dict 
    3762 
    3863def get_coords(file, delim='\t', skip=['#', '\n']): 
    3964    """ 
    4065    Return a list of x,y,z tuples from the coordinates listed in file. Lines 
    4166    starting with # and empty lines are ignored. 
     
    5378        if any(line.startswith(i) for i in skip): 
    5479            continue 
    5580        xyz.append(tuple(float(i) for i in line.split(delim))) 
    5681    return xyz 
    5782 
    58 def dist_sort(xyz, keyfunc=None): 
     83def calc_dist1(xyz, xyz0=None): 
    5984    """ 
     85    Add a distance column to the input coordinate list. 
     86 
    6087    xyz : list 
    6188        List of x,y,z tuples. 
    62     keyfunc : function 
    63         A function that returns a key to sort by, given an x,y,z,r tuple (r is 
    64         distance from x0,y0,z0), i.e. a comparator. 
     89    xyz0 : list 
     90        x,y,z location from which to calculated distance to all other points. 
    6591 
    6692    """ 
    67     def default_keyfunc(xyzr): 
    68         return xyzr[-1] 
    69  
    70     if keyfunc is None: 
    71         keyfunc = default_keyfunc 
    72  
    73     x0, y0, z0 = xyz[0] 
     93    if xyz0 is None: 
     94        x0, y0, z0 = xyz[0] 
     95    else: 
     96        x0, y0, z0 = xyz0 
    7497 
    7598    xyzr = [] 
    7699    for x, y, z in xyz: 
    77         r = dist(x, y, z, x0, y0, z0) 
     100        r = ((x - x0)**2 + (y - y0)**2 + (z - z0)**2)**0.5 
    78101        xyzr.append((x, y, z, r)) 
    79102 
     103    return xyzr 
     104 
     105def calc_dist2(xyz, xyz0=None): 
     106    """ 
     107    numpy version of calc_dist1. 
     108 
     109    """ 
     110    xyz = np.array(xyz) 
     111 
     112    if xyz0 is None: 
     113        x0, y0, z0 = xyz[0] 
     114    else: 
     115        x0, y0, z0 = xyz0 
     116 
     117    r = ((xyz[:,0] - x0)**2 + (xyz[:,1] - y0)**2 + (xyz[:,2] - z0)**2)**0.5 
     118    return np.c_[xyz, r] 
     119 
     120def dist_sort1(xyzr): 
     121    """ 
     122    Sort points by distance using the built-in sorted function. 
     123 
     124    xyzr : list 
     125        List of x,y,z,r tuples. 
     126 
     127    """ 
     128    def keyfunc(xyzr): 
     129        return xyzr[-1] 
     130 
    80131    return sorted(xyzr, key=keyfunc) 
     132 
     133def dist_sort2(xyzr): 
     134    """ 
     135    Sort using argsort array method, then use fancy indexing to return the 
     136    sorted xyzr array. 
     137 
     138    """ 
     139    try: 
     140        idx = xyzr[:,3].argsort() 
     141    except AttributeError: 
     142        xyzr = np.array(xyzr) 
     143        idx = xyzr[:,3].argsort() 
     144    return xyzr[idx] 
    81145 
    82146def pretty_print(xyzr): 
    83147    """ 
    84148    Print the contents of xyzr, a list of 4-tuples. 
    85149 
    86150    """ 
    87151    for x, y, z, r in xyzr: 
    88152        print '%.4f\t%.4f\t%.4f\t%.4f' % (x, y, z, r) 
    89153 
    90154def main(): 
     155    """ 
     156    Command line options 
     157    -------------------- 
     158    -xyz <x> <y> <z> : 
     159        Coordinates of the desired reference point for distance calculations. 
     160    -d <distfunc>: 
     161        Choose which function in func_dict to use when opt_dict['distfunc'] is 
     162        called. 
     163    -s <sortfunc>: 
     164        Choose which function in func_dict to use when opt_dict['sortfunc'] is 
     165        called. 
     166 
     167    """ 
    91168    filename = sys.argv[1] 
     169    opt_dict = set_options(sys.argv[1:], func_dict) 
    92170 
    93171    file = open(filename, 'r') 
    94172    xyz = get_coords(file) 
    95173    file.close() 
    96174 
    97     xyzr = dist_sort(xyz) 
     175    xyzr = opt_dict['distfunc'](xyz, opt_dict['xyz0']) 
     176    xyzr = opt_dict['sortfunc'](xyzr) 
    98177    pretty_print(xyzr) 
     178 
     179func_dict = {'calc_dist1': calc_dist1, 'calc_dist2': calc_dist2, 
     180             'dist_sort1': dist_sort1, 'dist_sort2': dist_sort2} 
    99181 
    100182if __name__ == '__main__': 
    101183    main() 
    102184 
    103 ### todo: 
    104 ### 1) try writing my own sorting function 
    105 ### 2) try redoing everything using numpy 
    106