Easily highlight source code for your blog with our Syntax Highlighter.

# Stou's I/O and sorting assignment

Unmodified
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