Scheduler that returns chunks of data to be queries on the K-D Tree. The number of chunks is determined by the number of processes.
| 102 | return _d.copy(), _i.astype(int).copy() |
| 103 | |
| 104 | class Scheduler: |
| 105 | """ |
| 106 | Scheduler that returns chunks of data to be queries on the K-D Tree. |
| 107 | The number of chunks is determined by the number of processes. |
| 108 | """ |
| 109 | def __init__(self, ndata, nprocs): |
| 110 | self._ndata = mp.RawValue(ctypes.c_int, ndata) |
| 111 | self._start = mp.RawValue(ctypes.c_int, 0) |
| 112 | self._lock = mp.Lock() |
| 113 | min_chunk = ndata // nprocs |
| 114 | min_chunk = ndata if min_chunk <= 2 else min_chunk |
| 115 | self._chunk = min_chunk |
| 116 | |
| 117 | def __iter__(self): |
| 118 | return self |
| 119 | |
| 120 | def next(self): # Python 2 support |
| 121 | self._lock.acquire() |
| 122 | ndata = self._ndata.value |
| 123 | start = self._start.value |
| 124 | chunk = self._chunk |
| 125 | if ndata: |
| 126 | if chunk > ndata: |
| 127 | _s0 = start |
| 128 | _s1 = start + ndata |
| 129 | self._ndata.value = 0 |
| 130 | else: |
| 131 | _s0 = start |
| 132 | _s1 = start + chunk |
| 133 | self._ndata.value = ndata - chunk |
| 134 | self._start.value = start + chunk |
| 135 | self._lock.release() |
| 136 | return slice(_s0, _s1) |
| 137 | else: |
| 138 | self._lock.release() |
| 139 | raise StopIteration |
| 140 | |
| 141 | def __next__(self): # Python 3 support |
| 142 | self._lock.acquire() |
| 143 | ndata = self._ndata.value |
| 144 | start = self._start.value |
| 145 | chunk = self._chunk |
| 146 | if ndata: |
| 147 | if chunk > ndata: |
| 148 | _s0 = start |
| 149 | _s1 = start + ndata |
| 150 | self._ndata.value = 0 |
| 151 | else: |
| 152 | _s0 = start |
| 153 | _s1 = start + chunk |
| 154 | self._ndata.value = ndata - chunk |
| 155 | self._start.value = start + chunk |
| 156 | self._lock.release() |
| 157 | return slice(_s0, _s1) |
| 158 | else: |
| 159 | self._lock.release() |
| 160 | raise StopIteration |