from bamt.utils.MathUtils import get_brave_matrix, get_proximity_matrix
[docs]
class BigBraveBN:
def __init__(self, n_nearest=5, threshold=0.3, proximity_metric="MI"):
self.n_nearest = n_nearest
self.threshold = threshold
self.proximity_metric = proximity_metric
self.possible_edges = []
[docs]
def set_possible_edges_by_brave(self, df):
"""Returns list of possible edges for structure learning
Args:
df (DataFrame): data
proximity_matrix (DataFrame): might be generated by get_mutual_info_score_matrix() function,
correlation etc.
n_nearest (int, optional): Number of Nearest neighbors, hyperparameter. Defaults to 5.
threshold (float, optional): Threshold for edge candidates to be passed to possible edges, threshold.
Defaults to 0.3 [0;1].
Returns:
Possible edges: list of possible edges
"""
proximity_matrix = get_proximity_matrix(
df, proximity_metric=self.proximity_metric
)
brave_matrix = get_brave_matrix(df.columns, proximity_matrix, self.n_nearest)
possible_edges_list = []
for c1 in df.columns:
for c2 in df.columns:
if (
brave_matrix.loc[c1, c2]
> brave_matrix.max(numeric_only="true").max() * self.threshold
):
possible_edges_list.append((c1, c2))
self.possible_edges = possible_edges_list