##// END OF EJS Templates
allow map objects to partition specified lengths
MinRK -
Show More
@@ -1,170 +1,167
1 1 # encoding: utf-8
2 2
3 3 """Classes used in scattering and gathering sequences.
4 4
5 5 Scattering consists of partitioning a sequence and sending the various
6 6 pieces to individual nodes in a cluster.
7 7
8 8
9 9 Authors:
10 10
11 11 * Brian Granger
12 12 * MinRK
13 13
14 14 """
15 15
16 16 #-------------------------------------------------------------------------------
17 17 # Copyright (C) 2008-2011 The IPython Development Team
18 18 #
19 19 # Distributed under the terms of the BSD License. The full license is in
20 20 # the file COPYING, distributed as part of this software.
21 21 #-------------------------------------------------------------------------------
22 22
23 23 #-------------------------------------------------------------------------------
24 24 # Imports
25 25 #-------------------------------------------------------------------------------
26 26
27 27 from __future__ import division
28 28
29 29 import types
30 30 from itertools import islice
31 31
32 32 from IPython.utils.data import flatten as utils_flatten
33 33
34 34 #-------------------------------------------------------------------------------
35 35 # Figure out which array packages are present and their array types
36 36 #-------------------------------------------------------------------------------
37 37
38 38 arrayModules = []
39 39 try:
40 40 import Numeric
41 41 except ImportError:
42 42 pass
43 43 else:
44 44 arrayModules.append({'module':Numeric, 'type':Numeric.arraytype})
45 45 try:
46 46 import numpy
47 47 except ImportError:
48 48 pass
49 49 else:
50 50 arrayModules.append({'module':numpy, 'type':numpy.ndarray})
51 51 try:
52 52 import numarray
53 53 except ImportError:
54 54 pass
55 55 else:
56 56 arrayModules.append({'module':numarray,
57 57 'type':numarray.numarraycore.NumArray})
58 58
59 59 class Map(object):
60 60 """A class for partitioning a sequence using a map."""
61 61
62 def getPartition(self, seq, p, q):
63 """Returns the pth partition of q partitions of seq."""
62 def getPartition(self, seq, p, q, n=None):
63 """Returns the pth partition of q partitions of seq.
64 64
65 The length can be specified as `n`,
66 otherwise it is the value of `len(seq)`
67 """
68 n = len(seq) if n is None else n
65 69 # Test for error conditions here
66 70 if p<0 or p>=q:
67 print "No partition exists."
68 return
71 raise ValueError("must have 0 <= p <= q, but have p=%s,q=%s" % (p, q))
69 72
70 N = len(seq)
71 remainder = N % q
72 basesize = N // q
73 remainder = n % q
74 basesize = n // q
73 75
74 76 if p < remainder:
75 77 low = p * (basesize + 1)
76 78 high = low + basesize + 1
77 79 else:
78 80 low = p * basesize + remainder
79 81 high = low + basesize
80 82
81 83 try:
82 84 result = seq[low:high]
83 85 except TypeError:
84 86 # some objects (iterators) can't be sliced,
85 87 # use islice:
86 88 result = list(islice(seq, low, high))
87 89
88 90 return result
89 91
90 92 def joinPartitions(self, listOfPartitions):
91 93 return self.concatenate(listOfPartitions)
92 94
93 95 def concatenate(self, listOfPartitions):
94 96 testObject = listOfPartitions[0]
95 97 # First see if we have a known array type
96 98 for m in arrayModules:
97 99 #print m
98 100 if isinstance(testObject, m['type']):
99 101 return m['module'].concatenate(listOfPartitions)
100 102 # Next try for Python sequence types
101 103 if isinstance(testObject, (types.ListType, types.TupleType)):
102 104 return utils_flatten(listOfPartitions)
103 105 # If we have scalars, just return listOfPartitions
104 106 return listOfPartitions
105 107
106 108 class RoundRobinMap(Map):
107 """Partitions a sequence in a roun robin fashion.
109 """Partitions a sequence in a round robin fashion.
108 110
109 111 This currently does not work!
110 112 """
111 113
112 def getPartition(self, seq, p, q):
113 # if not isinstance(seq,(list,tuple)):
114 # raise NotImplementedError("cannot RR partition type %s"%type(seq))
115 return seq[p:len(seq):q]
116 #result = []
117 #for i in range(p,len(seq),q):
118 # result.append(seq[i])
119 #return result
114 def getPartition(self, seq, p, q, n=None):
115 n = len(seq) if n is None else n
116 return seq[p:n:q]
120 117
121 118 def joinPartitions(self, listOfPartitions):
122 119 testObject = listOfPartitions[0]
123 120 # First see if we have a known array type
124 121 for m in arrayModules:
125 122 #print m
126 123 if isinstance(testObject, m['type']):
127 124 return self.flatten_array(m['type'], listOfPartitions)
128 125 if isinstance(testObject, (types.ListType, types.TupleType)):
129 126 return self.flatten_list(listOfPartitions)
130 127 return listOfPartitions
131 128
132 129 def flatten_array(self, klass, listOfPartitions):
133 130 test = listOfPartitions[0]
134 131 shape = list(test.shape)
135 132 shape[0] = sum([ p.shape[0] for p in listOfPartitions])
136 133 A = klass(shape)
137 134 N = shape[0]
138 135 q = len(listOfPartitions)
139 136 for p,part in enumerate(listOfPartitions):
140 137 A[p:N:q] = part
141 138 return A
142 139
143 140 def flatten_list(self, listOfPartitions):
144 141 flat = []
145 142 for i in range(len(listOfPartitions[0])):
146 143 flat.extend([ part[i] for part in listOfPartitions if len(part) > i ])
147 144 return flat
148 145 #lengths = [len(x) for x in listOfPartitions]
149 146 #maxPartitionLength = len(listOfPartitions[0])
150 147 #numberOfPartitions = len(listOfPartitions)
151 148 #concat = self.concatenate(listOfPartitions)
152 149 #totalLength = len(concat)
153 150 #result = []
154 151 #for i in range(maxPartitionLength):
155 152 # result.append(concat[i:totalLength:maxPartitionLength])
156 153 # return self.concatenate(listOfPartitions)
157 154
158 155 def mappable(obj):
159 156 """return whether an object is mappable or not."""
160 157 if isinstance(obj, (tuple,list)):
161 158 return True
162 159 for m in arrayModules:
163 160 if isinstance(obj,m['type']):
164 161 return True
165 162 return False
166 163
167 164 dists = {'b':Map,'r':RoundRobinMap}
168 165
169 166
170 167
General Comments 0
You need to be logged in to leave comments. Login now