class StartsWith(LibFcn): name = prefix + "startswith" sig = Sigs([Sig([{"haystack": P.Array(P.Wildcard("A"))}, {"needle": P.Array(P.Wildcard("A"))}], P.Boolean()), Sig([{"haystack": P.Array(P.Wildcard("A"))}, {"needle": P.Wildcard("A")}], P.Boolean())]) errcodeBase = 15110 def __call__(self, state, scope, pos, paramTypes, haystack, needle): if isinstance(needle, (list, tuple)): return needle == haystack[:len(needle)] else: if len(haystack) == 0: return False else: return needle == haystack[0]
class Rotate(LibFcn): name = prefix + "rotate" sig = Sig([{"a": P.Array(P.Wildcard("A"))}, {"steps": P.Int()}], P.Array(P.Wildcard("A"))) errcodeBase = 15190 def __call__(self, state, scope, pos, paramTypes, a, steps): if steps < 0: raise PFARuntimeException("steps out of range", self.errcodeBase + 0, self.name, pos) elif steps == 0 or len(a) == 0: return a else: index = steps % len(a) left, right = a[:index], a[index:] return right + left
class SymDiff(LibFcn): name = prefix + "symdiff" sig = Sig([{"a": P.Array(P.Wildcard("A"))}, {"b": P.Array(P.Wildcard("A"))}], P.Array(P.Wildcard("A"))) errcodeBase = 15540 def __call__(self, state, scope, pos, paramTypes, a, b): out = [] for ai in a: if ai not in b: out.append(ai) for bi in b: if bi not in a: out.append(bi) return out
class Permutations(LibFcn): name = prefix + "permutations" sig = Sig([{"a": P.Array(P.Wildcard("A"))}], P.Array(P.Array(P.Wildcard("A")))) errcodeBase = 15790 def __call__(self, state, scope, pos, paramTypes, a): out = [] i = 0 for permutation in itertools.permutations(a): i += 1 if i % 1000 == 0: state.checkTime() out.append(list(permutation)) return out
class Intersection(LibFcn): name = prefix + "intersection" sig = Sig([{"a": P.Array(P.Wildcard("A"))}, {"b": P.Array(P.Wildcard("A"))}], P.Array(P.Wildcard("A"))) errcodeBase = 15520 def __call__(self, state, scope, pos, paramTypes, a, b): out = [] for ai in a: if ai in b: out.append(ai) for bi in b: if bi in a: out.append(bi) return out
class Det(LibFcn): name = prefix + "det" sig = Sigs([ Sig([{ "x": P.Array(P.Array(P.Double())) }], P.Double()), Sig([{ "x": P.Map(P.Map(P.Double())) }], P.Double()) ]) errcodeBase = 24090 def __call__(self, state, scope, pos, paramTypes, x): if isinstance(x, (list, tuple)) and all( isinstance(xi, (list, tuple)) for xi in x): rows = len(x) if rows < 1: raise PFARuntimeException("too few rows/cols", self.errcodeBase + 0, self.name, pos) cols = len(x[0]) if cols < 1: raise PFARuntimeException("too few rows/cols", self.errcodeBase + 0, self.name, pos) if raggedArray(x): raise PFARuntimeException("ragged columns", self.errcodeBase + 1, self.name, pos) if rows != cols: raise PFARuntimeException("non-square matrix", self.errcodeBase + 2, self.name, pos) if any( any(math.isnan(z) or math.isinf(z) for z in row) for row in x): return float("nan") else: return float(np().linalg.det(arraysToMatrix(x))) elif isinstance(x, dict) and all( isinstance(x[i], dict) for i in list(x.keys())): keys = list(rowKeys(x).union(colKeys(x))) if len(keys) < 1 or all(len(row) == 0 for row in list(x.values())): raise PFARuntimeException("too few rows/cols", self.errcodeBase + 0, self.name, pos) if any( any( math.isnan(z) or math.isinf(z) for z in list(row.values())) for row in list(x.values())): return float("nan") else: return float(np().linalg.det(mapsToMatrix(x, keys, keys)))
class SimpleEuclidean(LibFcn): name = prefix + "simpleEuclidean" sig = Sig([{ "x": P.Array(P.Double()) }, { "y": P.Array(P.Double()) }], P.Double()) errcodeBase = 28000 def __call__(self, state, scope, pos, paramTypes, x, y): if len(x) != len(y): raise PFARuntimeException("dimensions of vectors do not match", self.errcodeBase + 0, self.name, pos) return math.sqrt(sum(((xi - yi)**2 for xi, yi in zip(x, y))))
class Sub(LibFcn): name = prefix + "sub" sig = Sigs([ Sig([{ "x": P.Array(P.Double()) }, { "y": P.Array(P.Double()) }], P.Array(P.Double())), Sig([{ "x": P.Array(P.Array(P.Double())) }, { "y": P.Array(P.Array(P.Double())) }], P.Array(P.Array(P.Double()))), Sig([{ "x": P.Map(P.Double()) }, { "y": P.Map(P.Double()) }], P.Map(P.Double())), Sig([{ "x": P.Map(P.Map(P.Double())) }, { "y": P.Map(P.Map(P.Double())) }], P.Map(P.Map(P.Double()))) ]) errcodeBase = 24040 def __call__(self, state, scope, pos, paramTypes, x, y): if isinstance(x, (list, tuple)) and all(isinstance(xi, (list, tuple)) for xi in x) and \ isinstance(y, (list, tuple)) and all(isinstance(yi, (list, tuple)) for yi in y): if len(x) != len(y) or any( len(xi) != len(yi) for xi, yi in zip(x, y)): raise PFARuntimeException("misaligned matrices", self.errcodeBase + 0, self.name, pos) return [[xj - yj for xj, yj in zip(xi, yi)] for xi, yi in zip(x, y)] elif isinstance(x, (list, tuple)) and isinstance(y, (list, tuple)): if len(x) != len(y): raise PFARuntimeException("misaligned matrices", self.errcodeBase + 0, self.name, pos) return [xi - yi for xi, yi in zip(x, y)] elif isinstance(x, dict) and all(isinstance(x[i], dict) for i in x.keys()) and \ isinstance(y, dict) and all(isinstance(y[i], dict) for i in y.keys()): rows = rowKeys(x).union(rowKeys(y)) cols = colKeys(x).union(colKeys(y)) return dict( (i, dict((j, x.get(i, {}).get(j, 0.0) - y.get(i, {}).get(j, 0.0)) for j in cols)) for i in rows) else: rows = rowKeys(x).union(rowKeys(y)) return dict((i, x.get(i, 0.0) - y.get(i, 0.0)) for i in rows)
class Combinations(LibFcn): name = prefix + "combinations" sig = Sig([{"a": P.Array(P.Wildcard("A"))}, {"size": P.Int()}], P.Array(P.Array(P.Wildcard("A")))) errcodeBase = 15780 def __call__(self, state, scope, pos, paramTypes, a, size): if size < 1: raise PFARuntimeException("size < 1", self.errcodeBase + 0, self.name, pos) out = [] i = 0 for combination in itertools.combinations(a, size): i += 1 if i % 1000 == 0: state.checkTime() out.append(list(combination)) return out
class Index(LibFcn): name = prefix + "index" sig = Sigs([Sig([{"haystack": P.String()}, {"pattern": P.String()}], P.Array(P.Int())), Sig([{"haystack": P.Bytes()}, {"pattern": P.Bytes()}], P.Array(P.Int()))]) errcodeBase = 35000 def __call__(self, state, scope, pos, paramTypes, haystack, pattern): haystack, pattern, to = convert(haystack, pattern, paramTypes[0]) re = Regexer(haystack, pattern, self.errcodeBase + 0, self.name, pos) if re.search(0): region = re.getRegion() out = [region.beg[0], region.end[0]] else: out = [] re.free() return out
class LinearVariance(LibFcn): name = prefix + "linearVariance" sig = Sigs([Sig([{"datum": P.Array(P.Double())}, {"model": P.WildRecord("M", {"covar": P.Array(P.Array(P.Double()))})}], P.Double()), Sig([{"datum": P.Array(P.Double())}, {"model": P.WildRecord("M", {"covar": P.Array(P.Array(P.Array(P.Double())))})}], P.Array(P.Double())), Sig([{"datum": P.Map(P.Double())}, {"model": P.WildRecord("M", {"covar": P.Map(P.Map(P.Double()))})}], P.Double()), Sig([{"datum": P.Map(P.Double())}, {"model": P.WildRecord("M", {"covar": P.Map(P.Map(P.Map(P.Double())))})}], P.Map(P.Double()))]) errcodeBase = 31010 def __call__(self, state, scope, pos, paramTypes, datum, model): covarType = [x["type"] for x in paramTypes[1]["fields"] if x["name"] == "covar"][0] if covarType == {"type": "array", "items": {"type": "array", "items": "double"}}: # sig1 datum = datum + [1.0] covar = model["covar"] if len(datum) != len(covar) or any(len(datum) != len(x) for x in covar): raise PFARuntimeException("misaligned covariance", self.errcodeBase + 0, self.name, pos) x = np().matrix([datum]) C = np().matrix(covar) return float(x.dot(C.dot(x.T))[0][0]) elif covarType == {"type": "array", "items": {"type": "array", "items": {"type": "array", "items": "double"}}}: # sig2 datum = datum + [1.0] x = np().matrix([datum]) out = [] for covar in model["covar"]: if len(datum) != len(covar) or any(len(datum) != len(x) for x in covar): raise PFARuntimeException("misaligned covariance", self.errcodeBase + 0, self.name, pos) C = np().matrix(covar) out.append(float(x.dot(C.dot(x.T))[0][0])) return out elif covarType == {"type": "map", "values": {"type": "map", "values": "double"}}: # sig3 datum = dict(list(datum.items()) + [("", 1.0)]) covar = model["covar"] keys = list(set(list(datum.keys()) + sum([list(x.keys()) for x in list(covar.values())], []))) x = np().matrix([[datum.get(k, 0.0) for k in keys]]) C = np().matrix([[covar.get(i, {}).get(j, 0.0) for j in keys] for i in keys]) return float(x.dot(C.dot(x.T))[0][0]) else: datum = dict(list(datum.items()) + [("", 1.0)]) covar = model["covar"] keys = list(set(list(datum.keys()) + sum(flatten([[list(x.keys()) for x in list(row.values())] for row in list(covar.values())]), []))) x = np().matrix([[datum.get(k, 0.0) for k in keys]]) out = {} for depkey, row in list(covar.items()): C = np().matrix([[row.get(i, {}).get(j, 0.0) for j in keys] for i in keys]) out[depkey] = float(x.dot(C.dot(x.T))[0][0]) return out
class NTile(LibFcn): name = prefix + "ntile" sig = Sig([{"a": P.Array(P.Wildcard("A"))}, {"p": P.Double()}], P.Wildcard("A")) errcodeBase = 15460 def __call__(self, state, scope, pos, paramTypes, a, p): if len(a) == 0: raise PFARuntimeException("empty array", self.errcodeBase + 0, self.name, pos) if math.isnan(p): raise PFARuntimeException("p not a number", self.errcodeBase + 1, self.name, pos) if p <= 0.0: return lowestN(a, 1, lambda x, y: compare(jsonNodeToAvroType(paramTypes[0]).items, x, y) < 0)[0] if p >= 1.0: return highestN(a, 1, lambda x, y: compare(jsonNodeToAvroType(paramTypes[0]).items, x, y) < 0)[0] sa = sorted(a, lambda x, y: compare(jsonNodeToAvroType(paramTypes[0]).items, x, y)) k = (len(a) - 1.0)*p f = math.floor(k) dataType = paramTypes[-1] if (dataType is "float") or (dataType is "double"): c = math.ceil(k) if f == c: return sa[int(k)] d0 = sa[int(f)] * (c - k) d1 = sa[int(c)] * (k - f) return d0 + d1 else: if len(sa) % 2 == 1: return sa[int(f)] else: return sa[int(k)]
class ClosestN(LibFcn): name = prefix + "closestN" sig = Sigs([ Sig([{"n": P.Int()}, {"datum": P.Array(P.Double())}, {"clusters": P.Array(P.WildRecord("C", {"center": P.Array(P.Double())}))}], P.Array(P.Wildcard("C"))), Sig([{"n": P.Int()}, {"datum": P.Wildcard("A")}, {"clusters": P.Array(P.WildRecord("C", {"center": P.Wildcard("B")}))}, {"metric": P.Fcn([P.Wildcard("A"), P.Wildcard("B")], P.Double())}], P.Array(P.Wildcard("C")))]) errcodeBase = 29010 def __call__(self, state, scope, pos, paramTypes, n, datum, clusters, *args): if n < 0: raise PFARuntimeException("n must be nonnegative", self.errcodeBase + 0, self.name, pos) if len(args) == 1: metric, = args distances = [callfcn(state, scope, metric, [datum, x["center"]]) for x in clusters] else: distances = [sum((di - xi)**2 for di, xi in zip(datum, x["center"])) for x in clusters] indexes = argLowestN(distances, n, lambda a, b: a < b) return [clusters[i] for i in indexes]
class Keys(LibFcn): name = prefix + "keys" sig = Sig([{"m": P.Map(P.Wildcard("A"))}], P.Array(P.String())) errcodeBase = 26010 def __call__(self, state, scope, pos, paramTypes, m): return m.keys()
class Values(LibFcn): name = prefix + "values" sig = Sig([{"m": P.Map(P.Wildcard("A"))}], P.Array(P.Wildcard("A"))) errcodeBase = 26020 def __call__(self, state, scope, pos, paramTypes, m): return m.values()
class Tanimoto(BinaryMetric): name = prefix + "tanimoto" sig = Sig([{ "x": P.Array(P.Boolean()) }, { "y": P.Array(P.Boolean()) }], P.Double()) errcodeBase = 28100 def __call__(self, state, scope, pos, paramTypes, x, y): length = len(x) if len(y) != length: raise PFARuntimeException("dimensions of vectors do not match", self.errcodeBase + 0, self.name, pos) a00, a01, a10, a11 = self.countPairs(x, y) return div((a11 + a00), (a11 + 2 * (a10 + a01) + a00))
class CompoundTest(LibFcn): name = prefix + "compoundTest" sig = Sig([{ "datum": P.WildRecord("D", {}) }, { "operator": P.String() }, { "comparisons": P.Array(P.WildRecord("T", {})) }, { "test": P.Fcn([P.Wildcard("D"), P.Wildcard("T")], P.Boolean()) }], P.Boolean()) errcodeBase = 32020 def __call__(self, state, scope, pos, paramTypes, datum, operator, comparisons, test): if operator == "and": for comparison in comparisons: if callfcn(state, scope, test, [datum, comparison]) is False: return False return True elif operator == "or": for comparison in comparisons: if callfcn(state, scope, test, [datum, comparison]) is True: return True return False elif operator == "xor": numTrue = 0 for comparison in comparisons: if callfcn(state, scope, test, [datum, comparison]) is True: numTrue += 1 return numTrue % 2 == 1 else: raise PFARuntimeException("unrecognized logical operator", self.errcodeBase + 0, self.name, pos)
class Scale(LibFcn): name = prefix + "scale" sig = Sigs([Sig([{"x": P.Array(P.Double())}, {"alpha": P.Double()}], P.Array(P.Double())), Sig([{"x": P.Array(P.Array(P.Double()))}, {"alpha": P.Double()}], P.Array(P.Array(P.Double()))), Sig([{"x": P.Map(P.Double())}, {"alpha": P.Double()}], P.Map(P.Double())), Sig([{"x": P.Map(P.Map(P.Double()))}, {"alpha": P.Double()}], P.Map(P.Map(P.Double())))]) errcodeBase = 24010 def __call__(self, state, scope, pos, paramTypes, x, alpha): if isinstance(x, (list, tuple)) and all(isinstance(xi, (list, tuple)) for xi in x): return [[xj * alpha for xj in xi] for xi in x] elif isinstance(x, (list, tuple)): return [xi * alpha for xi in x] elif isinstance(x, dict) and all(isinstance(x[i], dict) for i in x): return dict((i, dict((j, xj * alpha) for j, xj in xi.items())) for i, xi in x.items()) else: return dict((i, xi * alpha) for i, xi in x.items())
class FromSet(LibFcn, ObjKey): name = prefix + "fromset" sig = Sig([{"s": P.Map(P.Wildcard("A"))}], P.Array(P.Wildcard("A"))) errcodeBase = 26210 def __call__(self, state, scope, pos, paramTypes, s): return s.values()
class Split(LibFcn): name = prefix + "split" sig = Sig([{"m": P.Map(P.Wildcard("A"))}], P.Array(P.Map(P.Wildcard("A")))) errcodeBase = 26100 def __call__(self, state, scope, pos, paramTypes, m): return [{k: v} for k, v in m.items()]
class Sigmoid(LibFcn): name = prefix + "sigmoid" sig = Sig([{ "x": P.Array(P.Double()) }, { "y": P.Array(P.Double()) }, { "gamma": P.Double() }, { "intercept": P.Double() }], P.Double()) errcodeBase = 23030 def __call__(self, state, scope, pos, paramTypes, x, y, gamma, intercept): return math.tanh(gamma * dot(x, y, self.errcodeBase + 0, self.name, pos) + intercept)
class Len(LibFcn): name = prefix + "len" sig = Sig([{"a": P.Array(P.Wildcard("A"))}], P.Int()) errcodeBase = 15000 def genpy(self, paramTypes, args, pos): return "len({0})".format(*args) def __call__(self, state, scope, pos, paramTypes, a): return len(a)
class Subseq(LibFcn): name = prefix + "subseq" sig = Sig([{"a": P.Array(P.Wildcard("A"))}, {"start": P.Int()}, {"end": P.Int()}], P.Array(P.Wildcard("A"))) errcodeBase = 15010 def genpy(self, paramTypes, args, pos): return "{0}[{1}:{2}]".format(*args) def __call__(self, state, scope, pos, paramTypes, a, start, end): return a[start:end]
class Logit(LibFcn): name = prefix + "logit" sig = Sigs([ Sig([{ "x": P.Array(P.Double()) }], P.Array(P.Double())), Sig([{ "x": P.Map(P.Double()) }], P.Map(P.Double())), Sig([{ "x": P.Double() }], P.Double()) ]) errcodeBase = 25010 def __call__(self, state, scope, pos, paramTypes, x): return unwrapForNorm(x, lambda y: 1. / (1. + math.exp(-y)))
class Tanh(LibFcn): name = prefix + "tanh" sig = Sigs([ Sig([{ "x": P.Array(P.Double()) }], P.Array(P.Double())), Sig([{ "x": P.Map(P.Double()) }], P.Map(P.Double())), Sig([{ "x": P.Double() }], P.Double()) ]) errcodeBase = 25080 def __call__(self, state, scope, pos, paramTypes, x): return unwrapForNorm(x, lambda y: math.tanh(y))
class Softplus(LibFcn): name = prefix + "softplus" sig = Sigs([ Sig([{ "x": P.Array(P.Double()) }], P.Array(P.Double())), Sig([{ "x": P.Map(P.Double()) }], P.Map(P.Double())), Sig([{ "x": P.Double() }], P.Double()) ]) errcodeBase = 25060 def __call__(self, state, scope, pos, paramTypes, x): return unwrapForNorm(x, lambda y: math.log(1.0 + math.exp(y)))
class Cauchit(LibFcn): name = prefix + "cauchit" sig = Sigs([ Sig([{ "x": P.Array(P.Double()) }], P.Array(P.Double())), Sig([{ "x": P.Map(P.Double()) }], P.Map(P.Double())), Sig([{ "x": P.Double() }], P.Double()) ]) errcodeBase = 25050 def __call__(self, state, scope, pos, paramTypes, x): return unwrapForNorm(x, lambda y: 0.5 + (1. / math.pi) * math.atan(y))
class UpdateHoltWintersPeriodic(LibFcn): name = prefix + "updateHoltWintersPeriodic" sig = Sig([{ "x": P.Double() }, { "alpha": P.Double() }, { "beta": P.Double() }, { "gamma": P.Double() }, { "state": P.WildRecord( "A", { "level": P.Double(), "trend": P.Double(), "cycle": P.Array(P.Double()), "multiplicative": P.Boolean() }) }], P.Wildcard("A")) errcodeBase = 14050 def __call__(self, state, scope, pos, paramTypes, x, alpha, beta, gamma, theState): if alpha < 0.0 or alpha > 1.0: raise PFARuntimeException("alpha out of range", self.errcodeBase + 0, self.name, pos) if beta < 0.0 or beta > 1.0: raise PFARuntimeException("beta out of range", self.errcodeBase + 1, self.name, pos) if gamma < 0.0 or gamma > 1.0: raise PFARuntimeException("gamma out of range", self.errcodeBase + 2, self.name, pos) level_prev = theState["level"] trend_prev = theState["trend"] cycle_unrotated = theState["cycle"] if len(cycle_unrotated) == 0: raise PFARuntimeException("empty cycle", self.errcodeBase + 3, self.name, pos) cycle_rotated = cycle_unrotated[1:] + [cycle_unrotated[0]] cycle_prev = cycle_rotated[0] if theState["multiplicative"]: level = div(alpha * x, cycle_prev) + (1.0 - alpha) * (level_prev + trend_prev) trend = beta * (level - level_prev) + (1.0 - beta) * trend_prev cycle = div(gamma * x, level) + (1.0 - gamma) * cycle_prev else: level = alpha * (x - cycle_prev) + (1.0 - alpha) * (level_prev + trend_prev) trend = beta * (level - level_prev) + (1.0 - beta) * trend_prev cycle = gamma * (x - level) + (1.0 - gamma) * cycle_prev return dict(theState, level=level, trend=trend, cycle=([cycle] + cycle_rotated[1:]))
class EigenBasis(LibFcn): name = prefix + "eigenBasis" sig = Sigs([Sig([{"x": P.Array(P.Array(P.Double()))}], P.Array(P.Array(P.Double()))), Sig([{"x": P.Map(P.Map(P.Double()))}], P.Map(P.Map(P.Double())))]) errcodeBase = 24110 def calculate(self, x, size): symm = (x + x.T) * 0.5 evals, evects = np().linalg.eig(symm) evects = np().array(evects) evects2 = [evects[:,i] * (-1.0 if evects[0,i] < 0.0 else 1.0) for i in range(size)] eigvalm2 = [div(1.0, math.sqrt(abs(ei))) for ei in evals] order = np().argsort(eigvalm2) out = np().empty((size, size), dtype=np().double) for i in range(size): for j in range(size): out[i,j] = evects2[order[i]][j] * eigvalm2[order[i]] return out def __call__(self, state, scope, pos, paramTypes, x): if isinstance(x, (list, tuple)) and all(isinstance(xi, (list, tuple)) for xi in x): rows = len(x) if rows < 1: raise PFARuntimeException("too few rows/cols", self.errcodeBase + 0, self.name, pos) cols = len(x[0]) if cols < 1: raise PFARuntimeException("too few rows/cols", self.errcodeBase + 0, self.name, pos) if raggedArray(x): raise PFARuntimeException("ragged columns", self.errcodeBase + 1, self.name, pos) if rows != cols: raise PFARuntimeException("non-square matrix", self.errcodeBase + 2, self.name, pos) if any(any(math.isnan(z) or math.isinf(z) for z in row) for row in x): raise PFARuntimeException("non-finite matrix", self.errcodeBase + 3, self.name, pos) return matrixToArrays(self.calculate(arraysToMatrix(x), rows)) elif isinstance(x, dict) and all(isinstance(x[i], dict) for i in x.keys()): keys = list(rowKeys(x).union(colKeys(x))) if len(keys) < 1 or all(len(z) == 0 for z in x.values()): raise PFARuntimeException("too few rows/cols", self.errcodeBase + 0, self.name, pos) if any(any(math.isnan(z) or math.isinf(z) for z in row.values()) for row in x.values()): raise PFARuntimeException("non-finite matrix", self.errcodeBase + 3, self.name, pos) return matrixToMaps(self.calculate(mapsToMatrix(x, keys, keys), len(keys)), map(str, range(len(keys))), keys)
class BallR(LibFcn): name = prefix + "ballR" sig = Sigs([ Sig([{ "r": P.Double() }, { "datum": P.Array(P.Double()) }, { "codebook": P.Array(P.Array(P.Double())) }], P.Array(P.Array(P.Double()))), Sig([{ "r": P.Double() }, { "datum": P.Wildcard("A") }, { "codebook": P.Array(P.Wildcard("B")) }, { "metric": P.Fcn([P.Wildcard("A"), P.Wildcard("B")], P.Double()) }], P.Array(P.Wildcard("B"))) ]) errcodeBase = 30020 def __call__(self, state, scope, pos, paramTypes, r, datum, codebook, *args): if len(args) == 1: metric, = args distances = [ callfcn(state, scope, metric, [datum, x]) for x in codebook ] else: distances = [ math.sqrt(sum((di - xi)**2 for di, xi in zip(datum, x))) for x in codebook ] return [x for x, d in zip(codebook, distances) if d < r]