diff --git a/abducer/kb.py b/abducer/kb.py index 032731c..0ba1479 100644 --- a/abducer/kb.py +++ b/abducer/kb.py @@ -115,7 +115,80 @@ class add_KB(KBBase): def __len__(self): return sum(self._dict_len(v) for v in self.base.values()) +class ClsKB(KBBase): + def __init__(self, X, Y = None): + super().__init__() + self.base = {} + + if X is None: + return + if Y is None: + Y = [None] * len(X) + + for x, y in zip(X, Y): + self.base.setdefault(len(x), defaultdict(list))[y].append(np.array(x)) + + def get_candidates(self, key, length = None): + if key is None: + return self.get_all_candidates() + + length = self._length(length) + + return sum([self.base[l][key] for l in length], []) + + def get_all_candidates(self): + return sum([sum(v.values(), []) for v in self.base.values()], []) + + def _dict_len(self, dic): + return sum(len(c) for c in dic.values()) + + def __len__(self): + return sum(self._dict_len(v) for v in self.base.values()) + +class RegKB(KBBase): + def __init__(self, X, Y = None): + super().__init__() + tmp_dict = {} + for x, y in zip(X, Y): + tmp_dict.setdefault(len(x), defaultdict(list))[y].append(np.array(x)) + + self.base = {} + for l in tmp_dict.keys(): + data = sorted(list(zip(tmp_dict[l].keys(), tmp_dict[l].values()))) + X = [x for y, x in data] + Y = [y for y, x in data] + self.base[l] = (X, Y) + + def get_candidates(self, key, length = None): + if key is None: + return self.get_all_candidates() + + length = self._length(length) + + min_err = 999999 + candidates = [] + for l in length: + X, Y = self.base[l] + + idx = bisect.bisect_left(Y, key) + begin = max(0, idx - 1) + end = min(idx + 2, len(X)) + + for idx in range(begin, end): + err = abs(Y[idx] - key) + if abs(err - min_err) < 1e-9: + candidates.extend(X[idx]) + elif err < min_err: + candidates = copy.deepcopy(X[idx]) + min_err = err + return candidates + + def get_all_candidates(self): + return sum([sum(D[0], []) for D in self.base.values()], []) + + def __len__(self): + return sum([sum(len(x) for x in D[0]) for D in self.base.values()]) if __name__ == "__main__": pseudo_label_list = list(range(10)) @@ -130,4 +203,28 @@ if __name__ == "__main__": print() res = kb.get_candidates(7, length = 3) print(res) + print() + + X = ["1+1", "0+1", "1+0", "2+0", "1+0+1"] + Y = [2, 1, 1, 2, 2] + kb = ClsKB(X, Y) + print(len(kb)) + res = kb.get_candidates(2, 5) + print(res) + res = kb.get_candidates(2, 3) + print(res) + res = kb.get_candidates(None) + print(res) + print() + + X = ["1+1", "0+1", "1+0", "2+0", "1+0.5", "0.75+0.75"] + Y = [2, 1, 1, 2, 1.5, 1.5] + kb = RegKB(X, Y) + print(len(kb)) + res = kb.get_candidates(1.6) + print(res) + res = kb.get_candidates(1.6, length = 9) + print(res) + res = kb.get_candidates(None) + print(res)