| 106 | |
| 107 | |
| 108 | class PQEncoder(object): |
| 109 | |
| 110 | def __init__(self, dataset, code_bits=-1, bits_per_subvect=-1, |
| 111 | nsubvects=-1, elemwise_dist_func=dists_elemwise_sq): |
| 112 | X = dataset.X_train |
| 113 | self.elemwise_dist_func = elemwise_dist_func |
| 114 | |
| 115 | tmp = _parse_codebook_params(X.shape[1], code_bits=code_bits, |
| 116 | bits_per_subvect=bits_per_subvect, |
| 117 | nsubvects=nsubvects) |
| 118 | self.nsubvects, self.ncentroids, self.subvect_len = tmp |
| 119 | self.code_bits = int(np.log2(self.ncentroids)) |
| 120 | |
| 121 | # for fast lookups via indexing into flattened array |
| 122 | self.offsets = np.arange(self.nsubvects, dtype=np.int) * self.ncentroids |
| 123 | |
| 124 | self.centroids = _learn_centroids(X, self.ncentroids, self.nsubvects, |
| 125 | self.subvect_len) |
| 126 | |
| 127 | def name(self): |
| 128 | return "PQ_{}x{}b".format(self.nsubvects, self.code_bits) |
| 129 | |
| 130 | def params(self): |
| 131 | return {'_algo': 'PQ', '_ncodebooks': self.nsubvects, |
| 132 | '_code_bits': self.code_bits} |
| 133 | |
| 134 | def encode_X(self, X, **sink): |
| 135 | idxs = pq._encode_X_pq(X, codebooks=self.centroids) |
| 136 | return idxs + self.offsets # offsets let us index into raveled dists |
| 137 | |
| 138 | def encode_q(self, q, **sink): |
| 139 | return None # we use fit_query() instead, so fail fast |
| 140 | |
| 141 | def dists_true(self, X, q): |
| 142 | return np.sum(self.elemwise_dist_func(X, q), axis=-1) |
| 143 | |
| 144 | def fit_query(self, q, **sink): |
| 145 | self.q_dists_ = _fit_pq_lut(q, centroids=self.centroids, |
| 146 | elemwise_dist_func=self.elemwise_dist_func) |
| 147 | |
| 148 | def dists_enc(self, X_enc, q_unused=None): |
| 149 | # this line has each element of X_enc index into the flattened |
| 150 | # version of q's distances to the centroids; we had to add |
| 151 | # offsets to each col of X_enc above for this to work |
| 152 | centroid_dists = self.q_dists_.T.ravel()[X_enc.ravel()] |
| 153 | return np.sum(centroid_dists.reshape(X_enc.shape), axis=-1) |
| 154 | |
| 155 | |
| 156 | def _learn_best_quantization(luts): # luts can be a bunch of vstacked luts |
no outgoing calls
no test coverage detected