From 4ba174d50f291e5c9de71e7748c76832a7b6528d Mon Sep 17 00:00:00 2001 From: troyyyyy Date: Fri, 31 Mar 2023 21:32:34 +0800 Subject: [PATCH] [ENH] Add docstring for abducer_base --- abl/abducer/abducer_base.py | 177 ++++++++++++++++++++++++++++++------ 1 file changed, 151 insertions(+), 26 deletions(-) diff --git a/abl/abducer/abducer_base.py b/abl/abducer/abducer_base.py index 441223f..f300bfe 100644 --- a/abl/abducer/abducer_base.py +++ b/abl/abducer/abducer_base.py @@ -1,16 +1,3 @@ -# coding: utf-8 -# ================================================================# -# Copyright (C) 2021 Freecss All rights reserved. -# -# File Name ：abducer_base.py -# Author ：freecss -# Email ：karlfreecss@gmail.com -# Created Date ：2021/06/03 -# Description ： -# -# ================================================================# - -import os import abc import numpy as np from multiprocessing import Pool @@ -27,6 +14,23 @@ class AbducerBase(abc.ABC): self.mapping = dict(zip(self.kb.pseudo_label_list, list(range(len(self.kb.pseudo_label_list))))) def _get_cost_list(self, pred_res, pred_res_prob, candidates): + """ + Get the cost list of candidates based on the distance function. + + Parameters + ---------- + pred_res : list + The predicted result. + pred_res_prob : list + The predicted result probability. + candidates : list + The list of candidates. + + Returns + ------- + list + The cost list of candidates. + """ if self.dist_func == 'hamming': return hamming_dist(pred_res, candidates) @@ -35,6 +39,23 @@ class AbducerBase(abc.ABC): return confidence_dist(pred_res_prob, candidates) def _get_one_candidate(self, pred_res, pred_res_prob, candidates): + """ + Get the best candidate based on the distance function. + + Parameters + ---------- + pred_res : list + The predicted result. + pred_res_prob : list + The predicted result probability. + candidates : list + The list of candidates. + + Returns + ------- + list + The best candidate. + """ if len(candidates) == 0: return [] elif len(candidates) == 1 or self.zoopt: @@ -54,6 +75,25 @@ class AbducerBase(abc.ABC): return len(pred_res) def _zoopt_address_score(self, pred_res, pred_res_prob, key, sol): + """ + Get the address score for a single solution. + + Parameters + ---------- + sol_x : array-like + Solution to evaluate. + pred_res : list + List of predicted results. + pred_res_prob : list + List of probabilities for predicted results. + key : str + Key for the predicted results. + + Returns + ------- + float + The address score for the given solution. + """ address_idx = np.where(sol.get_x() != 0)[0] candidates = self.address_by_idx(pred_res, key, address_idx) if len(candidates) > 0: @@ -66,10 +106,28 @@ class AbducerBase(abc.ABC): return max_address_num - x.sum() def zoopt_get_solution(self, pred_res, pred_res_prob, key, max_address_num): + """Get the optimal solution using the Zoopt library. + + Parameters + ---------- + pred_res : list + List of predicted results. + pred_res_prob : list + List of probabilities for predicted results. + key : str + Key for the predicted results. + max_address_num : int or float + Maximum number of addresses to use. If float, represents the fraction of total addresses to use. + + Returns + ------- + array-like + The optimal solution. + """ length = len(flatten(pred_res)) dimension = Dimension(size=length, regs=[[0, 1]] * length, tys=[False] * length) objective = Objective( - lambda sol: self._zoopt_address_score(pred_res, pred_res_prob, key, sol), + lambda sol: self.zoopt_address_score(pred_res, pred_res_prob, key, sol), dim=dimension, constraint=lambda sol: self._constrain_address_num(sol, max_address_num), ) @@ -78,9 +136,42 @@ class AbducerBase(abc.ABC): return solution def address_by_idx(self, pred_res, key, address_idx): + """Get the addresses corresponding to the given indices. + + Parameters + ---------- + pred_res : list + List of predicted results. + key : str + Key for the predicted results. + address_idx : array-like + Indices of the addresses to retrieve. + + Returns + ------- + list + The addresses corresponding to the given indices. + """ return self.kb.address_by_idx(pred_res, key, address_idx) def abduce(self, data, max_address=-1, require_more_address=0): + """Perform abduction on the given data. + + Parameters + ---------- + data : tuple + Tuple containing the predicted results, predicted result probabilities, and key. + max_address : int or float, optional + Maximum number of addresses to use. If float, represents the fraction of total addresses to use. + If -1, use all addresses. Defaults to -1. + require_more_address : int, optional + Number of additional addresses to require. Defaults to 0. + + Returns + ------- + list + The abduced addresses. + """ pred_res, pred_res_prob, key = data assert(type(max_address) in (int, float)) @@ -103,17 +194,36 @@ class AbducerBase(abc.ABC): candidate = self._get_one_candidate(pred_res, pred_res_prob, candidates) return candidate - # def batch_abduce(self, Z, Y, max_address=-1, require_more_address=0): - # return [self.abduce((z, prob, y), max_address, require_more_address) for z, prob, y in zip(Z['cls'], Z['prob'], Y)] + def batch_abduce(self, Z, Y, max_address=-1, require_more_address=0): + """Perform abduction on the given data in batches. + + Parameters + ---------- + Z : list + List of predicted results. + Y : list + List of predicted result probabilities. + max_address : int or float, optional + Maximum number of addresses to use. If float, represents the fraction of total addresses to use. + If -1, use all addresses. Defaults to -1. + require_more_address : int, optional + Number of additional addresses to require. Defaults to 0. + + Returns + ------- + list + The abduced addresses. + """ + return [self.abduce((z, prob, y), max_address, require_more_address) for z, prob, y in zip(Z['cls'], Z['prob'], Y)] - def _batch_abduce_helper(self, args): - z, prob, y, max_address, require_more_address = args - return self.abduce((z, prob, y), max_address, require_more_address) + # def _batch_abduce_helper(self, args): + # z, prob, y, max_address, require_more_address = args + # return self.abduce((z, prob, y), max_address, require_more_address) - def batch_abduce(self, Z, Y, max_address=-1, require_more_address=0): - with Pool(processes=os.cpu_count()) as pool: - results = pool.map(self._batch_abduce_helper, [(z, prob, y, max_address, require_more_address) for z, prob, y in zip(Z['cls'], Z['prob'], Y)]) - return results + # def batch_abduce(self, Z, Y, max_address=-1, require_more_address=0): + # with Pool(processes=os.cpu_count()) as pool: + # results = pool.map(self._batch_abduce_helper, [(z, prob, y, max_address, require_more_address) for z, prob, y in zip(Z['cls'], Z['prob'], Y)]) + # return results def __call__(self, Z, Y, max_address=-1, require_more_address=0): return self.batch_abduce(Z, Y, max_address, require_more_address) @@ -134,7 +244,7 @@ class HED_Abducer(AbducerBase): candidate = self.address_by_idx(pred, k, address_idx) return candidate - def _zoopt_address_score(self, pred_res, pred_res_prob, key, sol): + def zoopt_address_score(self, pred_res, pred_res_prob, key, sol): all_address_flag = reform_idx(sol.get_x(), pred_res) lefted_idxs = [i for i in range(len(pred_res))] candidate_size = [] @@ -205,6 +315,21 @@ if __name__ == '__main__': print(res) print() + print('add_KB without GKB:, no cache') + kb = add_KB(use_cache=False) + abd = AbducerBase(kb, 'confidence') + res = abd.batch_abduce({'cls':[[1, 1]], 'prob':prob1}, [8], max_address=2, require_more_address=0) + print(res) + res = abd.batch_abduce({'cls':[[1, 1]], 'prob':prob2}, [8], max_address=2, require_more_address=0) + print(res) + res = abd.batch_abduce({'cls':[[1, 1]], 'prob':prob1}, [17], max_address=2, require_more_address=0) + print(res) + res = abd.batch_abduce({'cls':[[1, 1]], 'prob':prob1}, [17], max_address=1, require_more_address=0) + print(res) + res = abd.batch_abduce({'cls':[[1, 1]], 'prob':prob1}, [20], max_address=2, require_more_address=0) + print(res) + print() + print('prolog_KB with add.pl:') kb = prolog_KB(pseudo_label_list=list(range(10)), pl_file='../examples/datasets/mnist_add/add.pl') abd = AbducerBase(kb, 'confidence') @@ -241,9 +366,9 @@ if __name__ == '__main__': kb = add_KB() abd = AbducerBase(kb, 'confidence') - res = abd.batch_abduce({'cls':[[1, 1], [1, 2]], 'prob':multiple_prob}, [4, 8], max_address=4, require_more_address=0) + res = abd.batch_abduce({'cls':[[1, 1], [1, 2]], 'prob':multiple_prob}, [4, 8], max_address=2, require_more_address=0) print(res) - res = abd.batch_abduce({'cls':[[1, 1], [1, 2]], 'prob':multiple_prob}, [4, 8], max_address=4, require_more_address=1) + res = abd.batch_abduce({'cls':[[1, 1], [1, 2]], 'prob':multiple_prob}, [4, 8], max_address=2, require_more_address=1) print(res) print()