| 31 | return leaves |
| 32 | |
| 33 | class HierarchicalGrid: |
| 34 | # tree = {} |
| 35 | |
| 36 | # grids = [] |
| 37 | |
| 38 | def __init__(self, base, n_dims, n_layers, o, a, l): |
| 39 | cordo = [] |
| 40 | for _ in range(n_dims): |
| 41 | cordo.append(o+l/2) |
| 42 | self.root = Grid(-1, 0, n_layers, cordo, l) |
| 43 | self.base = base |
| 44 | self.n_dims = n_dims |
| 45 | self.n_layers = n_layers |
| 46 | self.o = o #座标最小值 |
| 47 | self.a = a #叶子节点长度 |
| 48 | self.l = l #总长度 |
| 49 | |
| 50 | # print("build hg with %d pivot, %d level. So all %d^(%d * %d) = %d grids" % ( |
| 51 | # n_dims, n_layers, base, n_dims, n_layers, base ** (n_layers * n_dims))) |
| 52 | # print("原点为%f，方格大小为%f" % (o, a)) |
| 53 | |
| 54 | def add_vector(self, vector, vec_id, emb): |
| 55 | # add one data |
| 56 | |
| 57 | now = self.root |
| 58 | pre = 0 |
| 59 | parts = self.base**self.n_dims |
| 60 | nowl = self.l |
| 61 | vec = copy.deepcopy(vector) |
| 62 | cord_now = copy.deepcopy(now.o) |
| 63 | |
| 64 | for i in range(self.n_layers): |
| 65 | bins = [] |
| 66 | nowl /= self.base |
| 67 | for j in range(self.n_dims): |
| 68 | bins.append(int(vec[j] / nowl)) |
| 69 | vec[j] %= nowl |
| 70 | cord_now[j] = vector[j]-vec[j]+nowl/2 |
| 71 | grid_id = self.parsing_grid_id(bins, pre, parts) |
| 72 | pre = grid_id |
| 73 | if grid_id in now.child: |
| 74 | grid = now.child[grid_id] |
| 75 | else: |
| 76 | grid = Grid(grid_id, i+1, self.n_layers, cord_now, nowl) |
| 77 | now.child[grid_id] = grid |
| 78 | if grid.is_leaf(): |
| 79 | grid.vector.append(vector) |
| 80 | grid.vec_ids.append(vec_id) |
| 81 | grid.emb.append(emb) |
| 82 | now = grid |
| 83 | # print("add vector in grid %d" % grid_id) |
| 84 | return now |
| 85 | |
| 86 | def parsing_grid_id(self, bins, pre, parts): |
| 87 | decimal = 0 |
| 88 | power = 0 |
| 89 | for i in range(len(bins)): |
| 90 | decimal += bins[i] * pow(self.base, power) |
no outgoing calls
no test coverage detected