text:
import networkx as nx
import numpy as np
import os
import pickle
import hdbscan
from dca.schemes import (
    DelaunayGraphParams,
    HDBSCANParams,
)
from hdbscan._hdbscan_linkage import label
import logging
from dca.loggers import logger_time
import gc

# -------------------------------------------------------------------------- #
# Logging settings
# -------------------------------------------------------------------------- #
logger = logging.getLogger("DCA_info_logger")
result_logger = logging.getLogger("DCA_result_logger")
time_logger = logging.getLogger("DCA_time_logger")
# -------------------------------------------------------------------------- #
# -------------------------------------------------------------------------- #

@logger_time
def _approximate_Delaunay_edges(
    root: str, input_array_filepath: str, graph_params: DelaunayGraphParams
):
    """
    Phase 1
    Approximates the Delaunay graph.
    :param root: root directory of the experiment.
    :param input_array_filepath: path to R and E input array.
    :param graph_params: Delaunay graph parameters.
    """
    # Load unfiltered Delaunay graph if exists
    unfiltered_Delaunay_edges_filepath = os.path.join(
        root, graph_params.unfiltered_edges_filepath
    )
    if not os.path.isfile(unfiltered_Delaunay_edges_filepath):
        os_output = os.system(
            "./dca/approximate_Delaunay_graph {0} --nrays {1} --out {2} --out_dist {3}".format(
                os.path.join(root, input_array_filepath),
                graph_params.T,
                unfiltered_Delaunay_edges_filepath,
                os.path.join(root, graph_params.unfiltered_edges_len_filepath),
            )
        )
        if os_output != 0:
            raise ValueError(
                "Executable file dca/approximate_Delaunay_graph not working. Forgotten to change ownership?"
            )

@logger_time
def _filter_Delaunay_edges(
    unfiltered_Delaunay_edges_filepath: str,
    unfiltered_Delaunay_edges_len_filepath: str,
    graph_params: DelaunayGraphParams,
    filtered_Delaunay_edges_filepath: str,
    filtered_Delaunay_edges_len_filepath: str,
    n_points: int,
    n_query_points: int = 0,
):
    """
    Filters the approximated Delaunay graph with given sphere coverage parameter B as
    explained in Section 3.2.
    :param unfiltered_Delaunay_edges_filepath: path to unfiltered approximated Delaunay graph.
    :param unfiltered_Delaunay_edges_len_filepath: path to unfiltered approximated Delaunay edge lengths.
    :param graph_params: Delaunay graph parameters.
    :param filtered_Delaunay_edges_filepath: path to filtered approximated Delaunay graph.
    :param filtered_Delaunay_edges_len_filepath: path to filtered approximated Delaunay edge lengths.
    :param n_points: total number of points in R and E.
    :param n_query_points: total number of query points in Q.
    :return: shape of the unfilitered Delaunay edges.
    """
    if not os.path.isfile(filtered_Delaunay_edges_filepath):
        unfiltered_Delaunay_edges = np.load(unfiltered_Delaunay_edges_filepath)
        unfiltered_Delaunay_edges_len = np.load(unfiltered_Delaunay_edges_len_filepath)
        logger.debug("Unfiltered edges loaded")
        logger.debug("Chosen sphere coverage: %s", graph_params.sphere_coverage)
        logger.debug(
            "Unfiltered Delaunay edges shape: %s", unfiltered_Delaunay_edges.shape
        )
        unfiltered_Delaunay_edges_len_shape = unfiltered_Delaunay_edges_len.shape

        vertex_stats = {i: [] for i in range(n_points + n_query_points)}
        for edge_idx in range(len(unfiltered_Delaunay_edges)):
            (v1, v2, s) = unfiltered_Delaunay_edges[edge_idx]
            d = unfiltered_Delaunay_edges_len[edge_idx]
            vertex_stats[v1].append([v2, s, d, edge_idx])
            vertex_stats[v2].append([v1, s, d, edge_idx])
        logger.debug("Vertex stats dict created")

        # Needed when working with large datasets
        if n_points > 20000:
            del unfiltered_Delaunay_edges
            del unfiltered_Delaunay_edges_len
            gc.collect()
            logger.debug("Unfiltered edges deleted")

        significant_edge_idxs = []
        for vtx_idx, vtx_data in vertex_stats.items():
            vtx_data = np.array(vtx_data)
            if len(vtx_data) != 0:
                vtx_data = vtx_data[(-vtx_data[:, 2]).argsort()]  # sort by distance
                vtx_data[:, 1] = np.divide(
                    vtx_data[:, 1], graph_params.T * 2
                )  # normalize significance
                idx = -1
                cur_coverage = 0
                while (
                    -idx <= len(vtx_data)
                    and cur_coverage < graph_params.sphere_coverage
                ):
                    significant_edge_idxs.append(int(vtx_data[idx, 3]))
                    cur_coverage += vtx_data[idx, 1]
                    idx += -1
        logger.debug("Significant edges extracted")

        # Needed when working with large datasets
        if n_points > 20000:
            del vertex_stats
            gc.collect()
            logger.debug("Vertex stats dict deleted")

        filtered_significant_edge_idxs, _ = np.unique(
            significant_edge_idxs, return_counts=True
        )
        logger.debug("Unique significant edges extracted")

        # Needed when working with large datasets
        if n_points > 20000:
            del significant_edge_idxs
            gc.collect()
            logger.debug("Duplicate significant edges deleted")
            unfiltered_Delaunay_edges = np.load(unfiltered_Delaunay_edges_filepath)
            unfiltered_Delaunay_edges_len = np.load(
                unfiltered_Delaunay_edges_len_filepath
            )
            logger.debug("Unfiltered edges loaded again")

        filtered_Delaunay_edges = unfiltered_Delaunay_edges[
            filtered_significant_edge_idxs, :2
        ]
        filtered_Delaunay_edges_len = unfiltered_Delaunay_edges_len[
            filtered_significant_edge_idxs
        ]
        logger.debug("Filtered Delaunay edges shape: %s", filtered_Delaunay_edges.shape)

        # Needed when working with large datasets
        if n_points > 20000:
            del unfiltered_Delaunay_edges
            del unfiltered_Delaunay_edges_len
            gc.collect()
            logger.debug("Unfiltered edges deleted")

        # Save filtered edges and their length
        np.save(filtered_Delaunay_edges_filepath, filtered_Delaunay_edges)
        np.save(filtered_Delaunay_edges_len_filepath, filtered_Delaunay_edges_len)
        return unfiltered_Delaunay_edges_len_shape

@logger_time
def _distil_Delaunay_connected_components(
    root: str, params: HDBSCANParams, graph: nx.Graph
):
    """
    Phase 2
    Distils the Delaunay graph into connected components using HDBSCAN as
    explained in Section 3.
    :param root: root directory of the experiment.
    :param params: HDBSCAN hyperparameters.
    :param graph: approximated Delaunay graph obtained in Phase 1.
    :return: array of labels corresponding to each point in R U E.
    """
    clusterer_path = os.path.join(root, params.clusterer_filepath)
    if os.path.isfile(clusterer_path):
        with open(clusterer_path, "rb") as f:
            clusterer = pickle.load(f)
    else:
        clusterer = hdbscan.HDBSCAN(
            min_cluster_size=params.min_cluster_size,
            min_samples=params.min_samples,
            metric="precomputed",
            cluster_selection_epsilon=params.cluster_selection_epsilon,
        )

        min_spanning_tree = nx.to_pandas_edgelist(
            nx.minimum_spanning_tree(graph)
        ).to_numpy()
        logger.debug("Minimum spanning tree extracted")

        del graph
        gc.collect()
        logger.debug("Graph deleted")

        min_spanning_tree = min_spanning_tree[np.argsort(min_spanning_tree.T[2]), :]

        # Convert edge list into standard hierarchical clustering format
        single_linkage_tree = label(min_spanning_tree)
        logger.debug("Single linkage tree extracted")

        (
            clusterer.labels_,
            clusterer.probabilities_,
            clusterer.cluster_persistence_,
            clusterer._condensed_tree,
            clusterer._single_linkage_tree,
        ) = hdbscan.hdbscan_._tree_to_labels(
            None,
            single_linkage_tree,
            params.min_cluster_size,
            cluster_selection_epsilon=params.cluster_selection_epsilon,
        )

        clusterer._min_spanning_tree = min_spanning_tree

        with open(clusterer_path, "wb") as f:
            pickle.dump(clusterer, f)
        logger.debug("HDBSCAN clusterer saved")

    labels = clusterer.labels_
    del clusterer
    gc.collect()

    return labels

@logger_time
def _sort_Delaunay_connected_components(
    root: str,
    input_data_labels: np.ndarray,
    clustering_params: HDBSCANParams,
    n_points: int,
):
    """
    Sorts the connected components obtained from HDBSCAN in decreasing order by size.
    :param root: root directory of the experiment.
    :param input_data_labels: labels of R and E points given by HDBSCAN.
    :param clustering_params: HDBSCAN hyperparameters.
    :param n_points: total number of points in R and E.
    :return: dict with keys representing component indices and arrays of
    corresponding vertices containined in each component as values;
    index of the first trivial component.
    """
    # Relabel outliers
    n_outliers = np.count_nonzero(input_data_labels == -1)
    per_outliers = n_outliers / n_points
    if per_outliers > 0.2:
        logger.warning(
            f"Warning: DCA obtained {per_outliers}% of outliers. Consider modifying HDBSCAN parameters."
        )
    outlier_idxs = np.arange(
        max(input_data_labels) + 1,
        max(input_data_labels) + 1 + n_outliers,
    )
    first_trivial_component_idx = max(input_data_labels) + 1
    input_data_labels[input_data_labels == -1] = outlier_idxs
    logger.debug("Outliers relabeled")
    logger.debug("Number of outliers: {0}".format(n_outliers))
    logger.debug(
        "Number of all connected components: {0}".format(
            n_outliers + first_trivial_component_idx
        )
    )

    # Relabel compoments by size
    unique_idx, counts = np.unique(input_data_labels, return_counts=True)
    relabel = np.argsort(counts.argsort()[::-1])[input_data_labels]
    np.save(
        os.path.join(root, clustering_params.input_array_labels_filepath),
        relabel,
    )
    logger.debug("Connected components relabeled by size")

    del input_data_labels
    gc.collect()
    logger.debug("Input data labels saved and removed from memory")

    component_vertex_idx_mapping = {idx: [] for idx in unique_idx}
    for label_idx, label in enumerate(relabel):
        component_vertex_idx_mapping[label].append(label_idx)
    logger.debug("Component vertex id mapping created")

    return component_vertex_idx_mapping, first_trivial_component_idx

def _extract_RE_subgraphs(graph: nx.Graph, num_R: int, num_E: int):
    """
    Extracts graph restrictions H^R and H^E from a given graph H.
    :param graph: Delaunay graph (or component) to extract restrictions from.
    :param num_R: total number of R points.
    :param num_E: total number of E points.
    :return: graph resticted to R; graph resticted to E.
    """
    R_graph = graph.subgraph(np.arange(num_R))
    E_graph = graph.subgraph(np.arange(num_R, num_R + num_E))
    return R_graph, E_graph

@logger_time
def _evaluate_Delaunay_component(
    subgraph_RE_comp: nx.Graph,
    subgraph_R_comp: nx.Graph,
    subgraph_E_comp: nx.Graph,
    num_R: int,
):
    """
    Phase 3
    Evaluates given graph-connected component.
    :param subgraph_RE_comp: Delaunay graph of the component containing R and E points..
    :param subgraph_R_comp: Delaunay graph of component restricted to R.
    :param subgraph_E_comp: Delaunay graph of component restricted to E.
    :param num_R: total number of R points.
    :return: indices of R points in the component; indices of E points in the component;
    component consistency score; component quality score; number of heteroeneous edges
    in the component; total number of edges in the component
    """

    comp_R_idxs = np.array(list(subgraph_R_comp.nodes()))
    comp_E_idxs = np.array(list(subgraph_E_comp.nodes())) - num_R
    comp_consistency = _get_graph_consistency(len(comp_R_idxs), len(comp_E_idxs))
    comp_quality, num_comp_RE_edges, num_total_comp_edges = _get_graph_quality(
        subgraph_RE_comp, subgraph_R_comp, subgraph_E_comp
    )
    return (
        comp_R_idxs,
        comp_E_idxs,
        comp_consistency,
        comp_quality,
        num_comp_RE_edges,
        num_total_comp_edges,
    )

@logger_time
def _get_graph_consistency(num_R_vertices: int, num_E_vertices: int):
    """
    Calculates the consistency of the given graph as in Definition 3.2.
    :param num_R_vertices: number of R vertices contained in a graph.
    :param num_E_vertices: number of E vertices contained in a graph.
    :return: consistency score of the graph.
    """
    return 1 - abs(num_R_vertices - num_E_vertices) / (num_R_vertices + num_E_vertices)

@logger_time
def _get_graph_quality(RE_graph: nx.Graph, R_graph: nx.Graph, E_graph: nx.Graph):
    """
    Calculates the quality of the given graph as in Definition 3.2.
    :param RE_graph: Delaunay graph (or component) containing both R and E points.
    :param R_graph: Delaunay graph (or component) restricted to R.
    :param E_graph: Delaunay graph (or component) restricted to E.
    :return: quality score of the graph; number of heterogeneous edges; total number of edges
    """
    num_R_and_E_edges = R_graph.number_of_edges() + E_graph.number_of_edges()
    num_total_edges = RE_graph.number_of_edges()
    num_RE_edges = num_total_edges - num_R_and_E_edges
    graph_quality = num_RE_edges / num_total_edges if num_total_edges != 0 else 0
    return graph_quality, num_RE_edges, num_total_edges

@logger_time
def _approximate_query_Delaunay_edges(
    exp_log_dir: str,
    input_array_filepath: str,
    query_input_array_filepath: str,
    graph_params: DelaunayGraphParams,
):
    """
    q-DCA
    Approximates edges to query points as in Secion 3.2,
    :param exp_log_dir: export directory to save results to.
    :param input_array_filepath: path to R and E input array.
    :param query_input_array_filepath: path to query points array.
    :param graph_params: Delaunay graph parameters.
    """
    query_Delaunay_edges_filepath = os.path.join(
        exp_log_dir, graph_params.query_unfiltered_edges_filename
    )
    # Load unfiltered Delaunay graph if exists
    if not os.path.isfile(query_Delaunay_edges_filepath):
        os_output = os.system(
            "./dca/approximate_Delaunay_graph {0} --queries {1} --nrays {2} --out {3} --out_dist {4}".format(
                input_array_filepath,
                query_input_array_filepath,
                graph_params.T * 2,
                query_Delaunay_edges_filepath,
                os.path.join(
                    exp_log_dir, graph_params.query_unfiltered_edges_len_filename
                ),
            )
        )
        if os_output != 0:
            raise ValueError(
                "Executable file dca/approximate_Delaunay_graph not working. Forgotten to change ownership?"
            )

max_stars_repo_name: petrapoklukar/DCA

max_issues_repo_name: petrapoklukar/DCA

max_forks_repo_name: petrapoklukar/DCA

reason: import numpy,import networkx


text:
import numpy as np
import pandas as pd

import pdb
import sys
import torch

from sklearn.metrics import average_precision_score, roc_auc_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from tqdm import tqdm
from itertools import product

from simulation_helpers import separable_decision_rule_1d, inseparable_decision_rule_1d
from simulation_helpers import separable_decision_rule_nd, inseparable_decision_rule_nd
from simulation_helpers import create_gap_1d, create_gap_nd
from simulation_helpers import generate_g, generate_y, generate_s_scar
from sarpu.PUmodels import LogisticRegressionPU
from km.Kernel_MPE_grad_threshold import wrapper as ramaswamy

from sarpu.pu_learning import *
from baselines import cdmm, supervised_rel_prior, sar_em_rel_prior, scar_km2_rel_prior
from method import train_relative_estimator, get_loss
from eval_fs import eval_relative_prior, eval_pred_prior
from gpu_utils import restrict_GPU_pytorch
restrict_GPU_pytorch(sys.argv[2])
import itertools
np.random.seed(42)

# Fixed experimental parameters
method = sys.argv[1]
n_runs = 5
n_groups = 2
seeds = [10, 52, 30, 42, 1000, 1, 2, 3, 4, 29, 82374, 9]
scar_assumption = True
classification_model_type = LogisticRegressionPU
classification_model = classification_model_type(fit_intercept=True)

results_f_name = 'results/robustness_' + method + '_underestimate_results'

c1 = .5
c2_opts = [.3]
separability_opts =  ['overlap']
group_gap_opts = [-2, -1, -.5, -.25, 0]
p_y_diff_opts = [0, .1, .2, .3]
lamda_opts = [0]
std, n_to_sample_g1, n_to_sample_g2, n_attributes = 16, 10000, 10000, 5

# Specifying which set of experiments to run

# Enumerating the experiment configurations
# Relevant experiment params = c1, c2, n_dims, std, n_examples, group_gap
# TODO: Add iterations over gap between decision rules?
# call it decision_rule_diff

all_combinations = list(product(c2_opts, separability_opts, group_gap_opts, p_y_diff_opts))

# Creating group configurations based on each experiment config
expmt_configs = []
for lamda in lamda_opts:
    for  c2,  separability_assumption, group_gap, p_y_diff in all_combinations:
        # configs for generating group features
        g1_mean = np.ones(n_attributes)
        g2_mean = g1_mean + group_gap

        g1_config = {'mean': g1_mean, 'std': std, 'n_samples': n_to_sample_g1,
                     'n_attributes': n_attributes, 'n_groups': n_groups, 'group_feat': 1}
        g2_config = {'mean': g2_mean, 'std': std, 'n_samples': n_to_sample_g2,
                     'n_attributes': n_attributes, 'n_groups': n_groups, 'group_feat': 2}

        expmt_config = {'scar_assumption': scar_assumption,
                        'separability_assumption': separability_assumption,
                        'labeling_frequency_g1': c1, 'labeling_frequency_g2': c2,
                        'group_gap': group_gap, 'lamda': lamda,
                       'n_groups': n_groups, 'n_attributes': n_attributes,
                       'p_y_diff': p_y_diff}
        expmt_configs.append((expmt_config, g1_config, g2_config))
# Iterating over all experiments
result_dicts = []
for expmt_config, g1_config, g2_config in tqdm(expmt_configs):
    # generate features
    x1 = generate_g(g1_config)
    x2 = generate_g(g2_config)

    # generate true labels
    # TODO: change to accommodate different decision rules for each group

    beta = 1
    y1 = generate_y(x1, g1_config, inseparable_decision_rule_nd, beta,
                    decision_rule_diff=expmt_config['p_y_diff'])
    y2 = generate_y(x2, g2_config, inseparable_decision_rule_nd, beta)
    # generate observed labels
    if expmt_config['scar_assumption']:
        s1 = generate_s_scar(y1, expmt_config['labeling_frequency_g1'])
        s2 = generate_s_scar(y2, expmt_config['labeling_frequency_g2'])

    x = np.concatenate([x1, x2])
    y = np.concatenate([y1, y2])
    s = np.concatenate([s1, s2])

    for run in range(n_runs):
        x1_train, x1_test, y1_train, y1_test, s1_train, s1_test = train_test_split(x1, y1, s1,  test_size=.4,random_state=seeds[run],shuffle=True, stratify=y1)
        x2_train, x2_test, y2_train, y2_test, s2_train, s2_test = train_test_split(x2, y2, s2, test_size=.4, random_state=seeds[run], shuffle=True, stratify=y2)

        x1_val, x1_test, y1_val, y1_test, s1_val, s1_test = train_test_split(x1_test, y1_test, s1_test,  test_size=.5, shuffle=True, stratify=y1_test)
        x2_val, x2_test, y2_val, y2_test, s2_val, s2_test = train_test_split(x2_test, y2_test, s2_test, test_size=.5, shuffle=True, stratify=y2_test)

        x_train = np.concatenate([x1_train, x2_train])
        y_train = np.concatenate([y1_train, y2_train])
        s_train = np.concatenate([s1_train, s2_train])

        x_test = np.concatenate([x1_test, x2_test])
        y_test = np.concatenate([y1_test, y2_test])
        s_test = np.concatenate([s2_test, s2_test])

        x_val = np.concatenate([x1_val, x2_val])
        s_val = np.concatenate([s1_val, s2_val])

        scale_factor = np.max(np.abs(x_train))
        x_train_norm = x_train/scale_factor
        x_val_norm = x_val/scale_factor
        x_test_norm = x_test/scale_factor

        g1_train_idxs = x_train[:,0] == 1
        g2_train_idxs = x_train[:,1] == 1
        g1_test_idxs = x_test[:,0] == 1
        g2_test_idxs = x_test[:,1] == 1

        info = {}
        if method == 'ours':
            f_model, losses = train_relative_estimator(x_train, s_train,
                                                       expmt_config, n_epochs=10000)

            pred_rel_prior, pred_g1_prior, pred_g2_prior = eval_relative_prior(x_test, f_model)

            s_preds = torch.squeeze(f_model(torch.Tensor(x_val).cuda()))
            val_loss = get_loss(s_val, s_preds)
            s_preds = torch.squeeze(s_preds).detach().cpu()
            auc = roc_auc_score(s_val, s_preds)
            auprc = average_precision_score(s_val, s_preds)
            info = {'auprc': auprc, 'auc': auc, 'val_loss': val_loss}

        elif method == 'supervised':
            results = supervised_rel_prior(x_train[g1_train_idxs],
                                           x_train[g2_train_idxs],
                                           y1_train, y2_train, x_test, g1_config,
                                           classification_model=classification_model)

            pred_rel_prior, pred_g1_prior, pred_g2_prior = results

        elif method == 'negative':
            results = supervised_rel_prior(x_train[g1_train_idxs],
                                           x_train[g2_train_idxs],
                                           s_train[g1_train_idxs],
                                           s_train[g2_train_idxs], x_test, g1_config,
                                           classification_model=classification_model)

            pred_rel_prior, pred_g1_prior, pred_g2_prior = results

        elif method == 'sar-em':
            results = sar_em_rel_prior(x_train_norm[g1_train_idxs],
                                       x_train_norm[g2_train_idxs],
                                       s_train[g1_train_idxs],
                                       s_train[g2_train_idxs], x_test_norm, g1_config,
                                       classification_model=classification_model)

            pred_rel_prior, pred_g1_prior, pred_g2_prior = results

        elif method == 'cdmm':
            results  = cdmm(x_train_norm[g1_train_idxs],
                            x_train_norm[g2_train_idxs],
                            s_train[g1_train_idxs], s_train[g2_train_idxs], g1_config)

            pred_rel_prior, pred_g1_prior, pred_g2_prior = results

        elif method == 'scar-km2':
            results = scar_km2_rel_prior(x1_train, x2_train, s1_train, s2_train)

            pred_rel_prior, pred_g1_prior, pred_g2_prior = results

        true_g1_prior = y_test[g1_test_idxs].mean()
        true_g2_prior = y_test[g2_test_idxs].mean()
        true_rel_prior = true_g1_prior / true_g2_prior

        result_dict = {'pred_rel_prior': pred_rel_prior, 'true_rel_prior': true_rel_prior,
                       'rel_prior_err': true_rel_prior - pred_rel_prior,
                       'pred_g1_prior': pred_g1_prior, 'pred_g2_prior': pred_g2_prior,
                       'true_g1_prior': true_g1_prior, 'true_g2_prior': true_g2_prior,
                       'method': method, 'run': run}
        result_dict.update(**{"g1_" + k: v for k, v in g1_config.items()})
        result_dict.update(**{"g2_" + k: v for k, v in g2_config.items()})
        result_dict.update(expmt_config)
        result_dict.update(info)
        result_dicts.append(result_dict)

    results_df = pd.DataFrame(result_dicts)
    results_df.to_csv(results_f_name)

ext: py

reason: import numpy


text:
from operator import itemgetter
import numpy as np
import pickle
from math import log

class _TreeNote:
    """
    This is a internal object that holds the current labeled feature and subtrees.

    """
    def __init__(self, value, index):
        self.sub_trees = {}
        self.value = value
        self.index = index

    def add_subtree(self, key, subtree):
        self.sub_trees[key] = subtree

class DecisionTree:
    """
    Decision Tree object handles basic operations
    """

    @staticmethod
    def load_tree(file_path):
        with open(file_path, 'rb') as file:
            return pickle.load(file)

    def __init__(self):
        self.root = None

    @staticmethod
    def _calculate_shannon_entropy(output):
        """
        Calculate shannon entropy: H = - sum(p(xi) * log2p(xi))
        :param sample: ndarray
        :return: entropy of the current array
        """
        num_entries = len(output)
        num_classes = {}
        for current in output:
            if current not in num_classes.keys():
                num_classes[current] = 1
            else:
                num_classes[current] += 1
        entropy = 0.0
        for key in num_classes:
            p = float(num_classes[key] / num_entries)
            entropy -= p * log(p, 2)
        return entropy

    @staticmethod
    def _split_data_set(sample, output, axis, value, feature_label=[]):
        if axis < 0 or axis > sample.shape[1]:
            raise ValueError("ERROR: Axis " + str(axis) + " is out of bounds: " + str(sample.shape[1]))
        new_sample = []
        new_output = []
        for i in range(sample.shape[0]):
            feature_vec = sample[i, :]
            if feature_vec[axis] == value:
                reduced_feature = feature_vec[0:axis]
                reduced_feature = np.append(reduced_feature, feature_vec[(axis+1):])
                new_sample.append(reduced_feature)
                new_output.append(output[i])
        sub_feature_labels = []
        if len(feature_label) > 0:
            sub_feature_labels = feature_label[0:axis]
            sub_feature_labels.extend(feature_label[axis+1:])
        return np.array(new_sample), new_output, sub_feature_labels

    def _select_feature_to_split(self, sample, output):
        num_samples, num_features = sample.shape
        initial_entropy = self._calculate_shannon_entropy(output)
        best_info_gain = 0
        feature_axis = -1
        for i in range(num_features):
            unique_values = np.unique(sample[:, i])
            entropy = 0.0
            for value in unique_values:
                data_set, split_output, _ = self._split_data_set(sample, output, i, value)
                prob = float(data_set.shape[0] / num_samples)
                entropy += prob * self._calculate_shannon_entropy(split_output)
            info_gain = initial_entropy - entropy
            if info_gain > best_info_gain:
                best_info_gain = info_gain
                feature_axis = i
        return feature_axis

    @staticmethod
    def _vote_count(output):
        class_cluster = {}
        for value in output:
            if value not in class_cluster.keys():
                class_cluster[value] = 1
            else:
                class_cluster[value] += 1
        sorted_cluster = sorted(class_cluster.items(), key=itemgetter(1), reverse=True)
        return sorted_cluster[0][0]

    def train(self, sample, output, store_path='', feature_label=[]):
        if sample.size == 0:
            raise ValueError("ERROR: Cannot use empty sample")
        if len(output) == 0:
            raise ValueError("ERROR: Cannot use empty output")
        if sample.shape[0] != len(output):
            raise ValueError("Error: The number of training sample must equal to the number of output")
        self._train(sample, output, feature_label=feature_label, node=self.root)
        if len(store_path) != 0:
            with open(store_path, 'wb') as output_file:
                pickler = pickle.Pickler(output_file, -1)
                pickler.dump(self)

    def predict(self, target):
        if len(target) == 0:
            raise ValueError('ERROR: Cannot use empty input')
        if self.root is None:
            raise RuntimeError("ERROR: tree hasn't been trained yet, train the tree first.")
        return self._predict(target, self.root)

    def _predict(self, target, node: _TreeNote):
        if len(node.sub_trees) == 0:
            return node.value
        feature_to_check = node.index
        feature_value = target[feature_to_check]

        sub_tree = node.sub_trees[feature_value]
        if not sub_tree:
            return sub_tree.name
        else:
            return self._predict(target, sub_tree)

    def _train(self, sample, output, current_feature_value=-1, feature_label=[], node=None):
        """
        Recursively generate decision tree. This is private and should not be used. Use train method instead.
        :param sample: numpy ndarray
        :param output: list
        :param current_feature_value: int, default 0
        :param feature_label: list
        :param node: current node
        """
        unique_output = np.unique(output)
        if unique_output.size == 1:
            if self.root is None:
                self.root = _TreeNote(output[0], 0)
            else:
                node.add_subtree(current_feature_value, _TreeNote(output[0], 0))
            return
        if sample.shape[1] == 0:
            value = self._vote_count(output)
            if self.root is None:
                self.root = _TreeNote(value, 0)
            else:
                node.add_subtree(current_feature_value, _TreeNote(value, 0))
            return

        feature_to_split = self._select_feature_to_split(sample, output)
        feature_name = feature_label[feature_to_split]
        sub_node = _TreeNote(feature_name, feature_to_split)
        if self.root is None:
            self.root = sub_node
        else:
            node.add_subtree(current_feature_value, sub_node)
        unique_feature_value = np.unique(sample[:, feature_to_split])

        for feature_value in unique_feature_value:
            sub_sample, sub_output, sub_feature_labels = self._split_data_set(sample, output, feature_to_split,
                                                                              feature_value, feature_label)
            self._train(sub_sample, sub_output, feature_value, sub_feature_labels, sub_node)

max_stars_repo_name: jasonhotsauce/neural-network

max_issues_repo_name: jasonhotsauce/neural-network

max_forks_repo_name: jasonhotsauce/neural-network


text:
import os
import numpy as np
import cv2
from numpy.linalg import norm
import heapq
import multiprocessing
import time
import logging

logging.basicConfig(format='%(asctime)s %(message)s', level=logging.DEBUG)

x_slice = 5
y_slice = 5

class FeatureExtractor(object):
    def __init__(self, group):
        self.group = group

    def prepare_reference(self, dataset, folder='Reference'):
        raise NotImplementedError()

    def extract_from_dataset(self):
        raise NotImplementedError()

    def clean_up(self):
        raise NotImplementedError()

all_group_list = [
    'dvd_covers',
    'cd_covers',
    'book_covers',
    'museum_paintings',
    'video_frames',
    'business_cards'
]

# * [Grid Color Moments](http://homepages.inf.ed.ac.uk/rbf/CVonline/LOCAL_COPIES/AV0405/KEEN/av_as2_nkeen.pdf)
class GridColorMomentsExtractor(FeatureExtractor):
    def __init__(self, group):
        super().__init__(group)
        self.prepare_reference()

    def prepare_reference(self):
        for group in all_group_list:
            if os.path.isfile(f'{group}-Reference.npy'):
                continue
            train_data_path = os.path.join('dataset', group)
            train_data_path = os.path.join(train_data_path, 'Reference')
            images = sorted([item for item in os.listdir(
                train_data_path) if item.endswith('.jpg')])
            total = len(images)

            logging.info('-'*30)
            logging.info(f'Preparing {group} reference...')
            logging.info('-'*30)

            image_moments = np.ndarray((total, x_slice, y_slice, 3, 3), np.float32)

            for idx, image_name in enumerate(images):
                img = cv2.imread(os.path.join(train_data_path, image_name))
                dx = img.shape[0] // x_slice
                dy = img.shape[1] // y_slice

                for i in range(x_slice):
                    for j in range(y_slice):
                        img_slice = img[i*dx:(i+1)*dx, j*dy:(j+1)*dy]
                        slice_moment = GridColorMomentsExtractor.color_moments(
                            img_slice)
                        image_moments[idx][i][j] = slice_moment

                if idx % 10 == 0:
                    logging.debug(f'Done: {idx}/{total} images')

            logging.info('Loading done.')

            np.save(f'{group}-Reference.npy', image_moments)
            logging.info(f'Saving to {group}-Reference.npy files done.')

    def extract_from_dataset(self):
        dataset_folder = os.path.join('dataset', self.group)
        subfolders = [item for item in os.listdir(
            dataset_folder) if item != "Reference" and not item.startswith(".")]
        top1_sum = 0
        top5_sum = 0
        total_count = 0
        for subfolder in subfolders:
            top1_count, top5_count, image_count = self._extract_from_subfolder(os.path.join(dataset_folder, subfolder))
            top1_sum += top1_count
            top5_sum += top5_count
            total_count += image_count

        return top1_sum, top5_sum, total_count

    def clean_up(self):
        for group in all_group_list:
            os.remove(f'{group}-Reference.npy')

    def _extract_from_subfolder(self, subfolder):
        images = sorted([item for item in os.listdir(
            subfolder) if item.endswith('.jpg')])

        total = len(images)

        logging.info('-'*30)
        logging.info(f'Classifying {subfolder}...')
        logging.info('-'*30)

        top1_count = 0
        top5_count = 0

        weight = np.reshape(np.array(x_slice * y_slice * [
            [1, 0, 0],
            [0, 2, 0],
            [0, 0, 1]
        ]), (x_slice, y_slice, 3, 3))

        for idx, image_name in enumerate(images):
            img = cv2.imread(os.path.join(subfolder, image_name))
            dx = img.shape[0] // x_slice
            dy = img.shape[1] // y_slice

            image_moment = np.ndarray((x_slice, y_slice, 3, 3), np.float32)

            for i in range(x_slice):
                for j in range(y_slice):
                    img_slice = img[i*dx:(i+1)*dx, j*dy:(j+1)*dy]
                    slice_moment = GridColorMomentsExtractor.color_moments(
                        img_slice)
                    image_moment[i][j] = slice_moment

            possible_values = []
            for group in all_group_list:
                reference = np.load(f'{group}-Reference.npy')
                for reference_index in range(reference.shape[0]):
                    substracted_mat = np.subtract(reference[reference_index], image_moment)
                    mat = np.matmul(substracted_mat, weight)
                    distance = norm(mat)
                    heapq.heappush(possible_values, (distance, reference_index, group))

            top5_values = []
            for i in range(5):
                top5_values.append(heapq.heappop(possible_values))

            if idx == top5_values[0][1] and self.group == top5_values[0][2]:
                top1_count += 1
                logging.debug(f'hit: {image_name}! Distance:{top5_values[0][0]}')

            for distance, reference_index, group in top5_values:
                if idx == reference_index and self.group == group:
                    top5_count += 1
                    logging.debug(f'hit in top 5: {image_name}! Distance:{distance}')

            if idx % 10 == 0:
                logging.debug(f'Done: {idx}/{total} images')

        return top1_count, top5_count, total

    @staticmethod
    def color_moments(img):
        if img is None:
            return
        # Convert BGR to HSV colorspace
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        # Split the channels - h,s,v
        h, s, v = cv2.split(hsv)
        # Initialize the color feature
        color_feature = []
        # N = h.shape[0] * h.shape[1]
        # The first central moment - average
        h_mean = np.mean(h)  # np.sum(h)/float(N)
        s_mean = np.mean(s)  # np.sum(s)/float(N)
        v_mean = np.mean(v)  # np.sum(v)/float(N)
        color_feature.append([h_mean, s_mean, v_mean])
        # The second central moment - standard deviation
        h_std = np.std(h)  # np.sqrt(np.mean(abs(h - h.mean())**2))
        s_std = np.std(s)  # np.sqrt(np.mean(abs(s - s.mean())**2))
        v_std = np.std(v)  # np.sqrt(np.mean(abs(v - v.mean())**2))
        color_feature.append([h_std, s_std, v_std])
        # The third central moment - the third root of the skewness
        h_skewness = np.mean(abs(h - h.mean())**3)
        s_skewness = np.mean(abs(s - s.mean())**3)
        v_skewness = np.mean(abs(v - v.mean())**3)
        h_thirdMoment = h_skewness**(1./3)
        s_thirdMoment = s_skewness**(1./3)
        v_thirdMoment = v_skewness**(1./3)
        color_feature.append([h_thirdMoment, s_thirdMoment, v_thirdMoment])

        return np.array(color_feature)

class GaborExtractor(FeatureExtractor):

    def __init__(self, group):
        super().__init__(group)
        self.kernels = GaborExtractor.gabor_kernel()
        self.prepare_reference()

    @staticmethod
    def convolution(image, filters):
        result = np.zeros_like(image, dtype=float)
        # Normalize images for better comparison.
        # image = (image - image.mean()) // image.std()
        for kern in filters:
            fimg = cv2.filter2D(image, cv2.CV_64FC3, kern)
            np.maximum(result, fimg, result)
        return result

    @staticmethod
    def gabor_kernel():
        filters = []
        ksize = 9
        # define the range for theta and nu
        for theta in np.arange(0, np.pi, np.pi / 8):
            for nu in np.arange(np.pi / 4, 6*np.pi/4, np.pi / 4):
                kern = cv2.getGaborKernel(
                    (ksize, ksize), 1.0, theta, nu, 0.5, 0, ktype=cv2.CV_32F)
                kern /= 1.5*kern.sum()
                filters.append(kern)
        return np.array(filters)

    def feature_from_image(self, image):
        # initializing the feature vector
        feat = []
        # calculating the local energy for each convolved image
        for j in range(40):
            res = GaborExtractor.convolution(image, self.kernels[j])
            feat.append(np.mean(res))
            feat.append(np.std(res))
        return np.array(feat)

    def prepare_reference(self):
        for group in all_group_list:
            if os.path.isfile(f'{group}-Reference.npy'):
                continue
            train_data_path = os.path.join('dataset', group)
            train_data_path = os.path.join(train_data_path, 'Reference')
            image_list = sorted([item for item in os.listdir(
                train_data_path) if item.endswith('.jpg')])
            total = len(image_list)

            logging.info('-'*30)
            logging.info(f'Preparing {group} reference...')
            logging.info('-'*30)

            features = np.ndarray((total, 80), np.float32)

            # async_results = []
            # pool = multiprocessing.Pool(processes=multiprocessing.cpu_count())
            for idx, image_name in enumerate(image_list):
                img = cv2.imread(os.path.join(train_data_path, image_name))
                feature, idx = self._reference_worker(img, idx)
                features[idx] = feature
                if idx % 10 == 0:
                    logging.debug(f'Done: {idx}/{total} images')
            #     async_results.append(pool.apply_async(
            #         self._reference_worker, (img, idx)))

            # for async_result in async_results:
            #     feature, idx = async_result.get()
            #     features[idx] = feature
            #     if idx % 10 == 0:
            #         logging.debug(f'Done: {idx}/{total} images')
            # pool.close()
            # pool.join()
            logging.info('Loading done.')

            np.save(f'{group}-Reference.npy', features)
            logging.info(f'Saving to {group}-Reference.npy files done.')

    def extract_from_dataset(self):
        dataset_folder = os.path.join('dataset', self.group)
        subfolders = [item for item in os.listdir(
            dataset_folder) if item != "Reference" and not item.startswith(".")]
        top1_sum = 0
        top5_sum = 0
        total_count = 0
        for subfolder in subfolders:
            top1_count, top5_count, image_count = self._extract_from_subfolder(os.path.join(dataset_folder, subfolder))
            top1_sum += top1_count
            top5_sum += top5_count
            total_count += image_count

        return top1_sum, top5_sum, total_count

    def clean_up(self):
        for group in all_group_list:
            os.remove(f'{group}-Reference.npy')

    def _extract_from_subfolder(self, subfolder):
        images = sorted([item for item in os.listdir(
            subfolder) if item.endswith('.jpg')])

        total = len(images)
        top1_count = 0
        top5_count = 0

        print('-'*30)
        print(f'Classifying {subfolder}...')
        print('-'*30)
        # pool = multiprocessing.Pool(processes=multiprocessing.cpu_count())
        # async_results = []
        for idx, image_name in enumerate(images):
            img = cv2.imread(os.path.join(subfolder, image_name))
            img = cv2.resize(img, (0, 0), fx=0.25, fy=0.25)
            top1, top5 = self._worker(idx, img, image_name, total)
            top1_count += top1
            top5_count += top5
            # async_results.append(pool.apply_async(
            #     self._worker, (idx, img, image_name, total)))

        # for async_result in async_results:
        #     top1, top5 = async_result.get()
        #     top1_count += top1
        #     top5_count += top5
        # pool.close()
        # pool.join()

        return top1_count, top5_count, total

    def _reference_worker(self, img, idx):
        feature = self.feature_from_image(img)
        return feature, idx

    def _worker(self, idx, img, image_name, total):
        feature = self.feature_from_image(img)
        possible_values = []
        top1_count = 0
        top5_count = 0
        for group in all_group_list:
            reference = np.load(f'{group}-Reference.npy')
            for reference_index in range(reference.shape[0]):
                substracted_mat = np.subtract(
                    reference[reference_index], feature)
                distance = norm(substracted_mat)
                heapq.heappush(possible_values, (distance, reference_index, group))
        top5_values = []
        for i in range(5):
            top5_values.append(heapq.heappop(possible_values))

        if idx == top5_values[0][1] and self.group == top5_values[0][2]:
            top1_count += 1
            logging.debug(f'hit: {image_name}! Distance:{top5_values[0][0]}')

        for distance, reference_index, group in top5_values:
            if idx == reference_index and self.group == group:
                top5_count += 1
                logging.debug(f'hit in top 5: {image_name}! Distance:{distance}')

        if idx % 10 == 0:
            logging.debug(f'Done: {idx}/{total} images')

        return top1_count, top5_count

class DogSIFTExtactor(FeatureExtractor):
    def __init__(self, group):
        super().__init__(group)
        self.prepare_reference()
        # if not os.path.isfile(f'{dataset_folder}-Reference.npy'):
        #     self.prepare_reference(dataset_folder)
        # test_list = np.load(f'{dataset_folder}-Reference.npy')
        # self.reference = []
        # # self.reference = []
        # for index in range(test_list.shape[0]):
        #     test_entry = test_list[index]
        #     keypoints = []
        #     for entry in test_entry[0]:
        #         keypoint = cv2.KeyPoint(
        #             x=entry[0],
        #             y=entry[1],
        #             _size=entry[2],
        #             _angle=entry[3],
        #             _response=entry[4],
        #             _octave=int(entry[5]),
        #             _class_id=int(entry[6]))
        #         keypoints.append(keypoint)
        #     self.reference.append((keypoints, test_entry[1]))

    def prepare_reference(self):
        for group in all_group_list:
            if os.path.isfile(f'{group}-Reference.npy'):
                continue
            train_data_path = os.path.join('dataset', group)
            train_data_path = os.path.join(train_data_path, 'Reference')
            images = sorted([item for item in os.listdir(
                train_data_path) if item.endswith('.jpg')])
            total = len(images)

            logging.info('-'*30)
            logging.info(f'Preparing {group} reference...')
            logging.info('-'*30)
            results = []
            for idx, image_name in enumerate(images):
                img = cv2.imread(os.path.join(train_data_path, image_name))
                results.append(self._reference_worker(img, idx)[1:])
                if idx % 10 == 0:
                    logging.debug(f'Done: {idx}/{total} images')

            logging.info('Loading done.')

            np.save(f'{group}-Reference.npy', np.array(results))
            logging.info(f'Saving to {group}-Reference.npy files done.')

    def _reference_worker(self, img, idx):
        sift = cv2.xfeatures2d.SIFT_create()
        keypoints, features = sift.detectAndCompute(img, None)
        result = []
        for index, keypoint in enumerate(keypoints):
            entry = [
                keypoint.pt[0],
                keypoint.pt[1],
                keypoint.size,
                keypoint.angle,
                keypoint.response,
                keypoint.octave,
                keypoint.class_id
            ]
            result.append(entry)
        return idx, np.array(result), features

    def extract_from_dataset(self):
        dataset_folder = os.path.join('dataset', self.group)
        subfolders = [item for item in os.listdir(
            dataset_folder) if item != "Reference" and not item.startswith(".")]
        top1_sum = 0
        top5_sum = 0
        total_count = 0
        for subfolder in subfolders:
            top1_count, top5_count, image_count = self._extract_from_subfolder(os.path.join(dataset_folder, subfolder))
            top1_sum += top1_count
            top5_sum += top5_count
            total_count += image_count

        return top1_sum, top5_sum, total_count

    def clean_up(self):
        for group in all_group_list:
            os.remove(f'{group}-Reference.npy')

    def _extract_from_subfolder(self, subfolder):
        images = sorted([item for item in os.listdir(
            subfolder) if item.endswith('.jpg')])

        total = len(images)
        top1_count = 0
        top5_count = 0

        logging.info('-'*30)
        logging.info(f'Classifying {subfolder}...')
        logging.info('-'*30)
        start = time.time()
        pool = multiprocessing.Pool(processes=multiprocessing.cpu_count())
        async_results = []
        for idx, image_name in enumerate(images):
            img = cv2.imread(os.path.join(subfolder, image_name))
            # top1, top5 = self._worker(idx, img, image_name, total)
            # top1_count += top1
            # top5_count += top5
            async_results.append(pool.apply_async(
                self._worker, (idx, img, image_name, total)))
        pool.close()
        pool.join()

        for async_result in async_results:
            top1, top5 = async_result.get()
            top1_count += top1
            top5_count += top5

        end = time.time()
        logging.debug(f"time:{end - start}")
        return top1_count, top5_count, total

    def _worker(self, idx, img, image_name, total):
        sift = cv2.xfeatures2d.SIFT_create()
        keypoints, features = sift.detectAndCompute(img, None)

        index_params = {'algorithm': 1, 'trees': 5}
        search_params = {'checks': 50}   # or pass empty dictionary
        flann = cv2.FlannBasedMatcher(index_params, search_params)

        possible_values = []
        top1_count = 0
        top5_count = 0
        for group in all_group_list:
            test_list = np.load(f'{group}-Reference.npy')
            reference = []
            for index in range(test_list.shape[0]):
                test_entry = test_list[index]
                keypoints = []
                for entry in test_entry[0]:
                    keypoint = cv2.KeyPoint(
                        x=entry[0],
                        y=entry[1],
                        _size=entry[2],
                        _angle=entry[3],
                        _response=entry[4],
                        _octave=int(entry[5]),
                        _class_id=int(entry[6]))
                    keypoints.append(keypoint)

                reference.append((keypoints, test_entry[1]))

            for reference_index, test_entry in enumerate(reference):
                matches = flann.knnMatch(test_entry[1], features, k=2)
                matches = [(m, n) for (m, n) in matches if m.distance < 0.7 * n.distance]
                heapq.heappush(possible_values, (-len(matches), reference_index, group))

        top5_values = []
        for i in range(5):
            top5_values.append(heapq.heappop(possible_values))

        if idx == top5_values[0][1] and self.group == top5_values[0][2]:
            top1_count += 1
            logging.debug(f'hit: {image_name}! Distance:{top5_values[0][0]}')

        for distance, reference_index, group in top5_values:
            if idx == reference_index and self.group == group:
                top5_count += 1
                logging.debug(f'hit in top 5: {image_name}! Distance:{distance}')

        if idx % 10 == 0:
            logging.debug(f'Done: {idx}/{total} images')

        return top1_count, top5_count

reason: import numpy,from numpy


text:
import os
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from uncertainty_library.mc_drop import get_mc_predictions
from uncertainty_library.data import read_format_data_CelebA, ds2numpy
from uncertainty_library.attribution_methods import (
    get_importances,
    get_importances_vanilla,
    get_importances_CLUE,
    get_importances_LIME,
    get_importances_SHAP,
)
from uncertainty_library.plotting_functions import plot_that_comparison

train_ds, valid_ds = read_format_data_CelebA('img_align_celeba', category='Smiling')
x_train, y_train = ds2numpy(train_ds, max_num_batches=200)
x_valid, y_valid = ds2numpy(valid_ds, max_num_batches=200)
idx2label = ['neg', 'pos']

# Load the models
encoder = tf.keras.models.load_model('saved_models/dfc-vae-cnn/encoder')
decoder = tf.keras.models.load_model('saved_models/dfc-vae-cnn/decoder')
my_model = tf.keras.models.load_model('saved_models/classifier_smiling')

# Get predictions across the test data
mc_predictions, summary = get_mc_predictions(my_model, (x_valid, y_valid))
summary_selected = summary.loc[summary.Label == summary.Pred]
summary_selected = summary_selected[
    (summary_selected.Entropy > 0.35) & (summary_selected.Entropy < 0.45)
]

# Choose image and produce importances
# =============================================================================

os.makedirs('plots_comparisons', exist_ok=True)
for idx in summary_selected.sort_values(by="Epistemic").index[:50]:
    x = x_valid[idx]
    y = int(summary_selected.loc[idx].Label)
    label = idx2label[y]

    # Importances
    ig_entr, ig_ale, ig_epi = get_importances(x, my_model, encoder, decoder,
                                              known_class=y, jacobian_splits=16)

    ig_entr_vanilla, ig_ale_vanilla, ig_epi_vanilla = get_importances_vanilla(x, my_model)

    entr_clue = get_importances_CLUE(x, my_model, encoder, decoder)

    entr_lime = get_importances_LIME(x, my_model, bs=4, kernel_size=3,
                                     max_dist=10, ratio=2.0)

    entr_shap = get_importances_SHAP(x, my_model, x_train, bs=4)

    # Plots
    fig = plot_that_comparison(x, ig_entr, ig_entr_vanilla, entr_clue,
                               entr_lime, entr_shap, label=label, alpha=0.5)
    fig.savefig(f'plots_comparisons/{idx}.png',
                bbox_inches='tight', pad_inches=0.1)
    plt.close(fig)

max_stars_repo_name: Featurespace/uncertainty-attribution

max_issues_repo_name: Featurespace/uncertainty-attribution

max_forks_repo_name: Featurespace/uncertainty-attribution


text:
# -*- coding: utf-8 -*-
# pylint: disable=redefined-outer-name
"""Tests for the `PhonopyPpCalculation` class."""
import numpy as np

from aiida import orm
from aiida_phonopy.calculations.functions.get_force_sets import generate_get_force_sets

def test_get_force_sets(generate_workchain_force_sets):
    """Return only those inputs that the parser will expect to be there."""
    # First we need a pre ran `ForceSetsWorkChain`
    process = generate_workchain_force_sets()
    process.setup()  # we run the phonopy preprocess

    # Generating force_sets
    forces_dict = {}
    supercells = process.ctx.supercells
    base_force = np.array([1, 0, 0])

    for key, value in supercells.items():  # key: e.g. "supercell_001", "phonon_supercell_001"
        num = key.split("_")[-1]  # e.g. "001"
        if num == "forces":
            num = 0

        fake_forces = []
        for atom in range(len(value.sites)):
            fake_forces.append(base_force)
            base_force += base_force
        fake_forces = np.array(fake_forces)

        forces = orm.ArrayData()
        forces.set_array("forces", fake_forces)

        forces_dict[f"forces_{num}"] = forces  # fake ArrayData

    get_force_sets = generate_get_force_sets()
    output = get_force_sets(**forces_dict)

    assert "force_sets" in output
    assert len(output["force_sets"].get_array("force_sets")) == 1
    assert isinstance(output["force_sets"], orm.ArrayData)

max_stars_repo_name: bastonero/aiida-phonopy

max_issues_repo_name: bastonero/aiida-phonopy

max_forks_repo_name: bastonero/aiida-phonopy


text:
#!/usr/bin/env python
import freenect
import cv2
import os
import numpy as np
import time

def pretty_depth(depth):
    """Converts depth into a 'nicer' format for display

    This is abstracted to allow for experimentation with normalization

    Args:
        depth: A numpy array with 2 bytes per pixel

    Returns:
        A numpy array that has been processed with unspecified datatype
    """
    np.clip(depth, 0, 2**10 - 1, depth)
    depth >>= 2
    depth = depth.astype(np.uint8)
    return depth

def pretty_depth_cv(depth):
    """Converts depth into a 'nicer' format for display

    This is abstracted to allow for experimentation with normalization

    Args:
        depth: A numpy array with 2 bytes per pixel

    Returns:
        A numpy array with unspecified datatype
    """
    return pretty_depth(depth)

def video_cv(video):
    """Converts video into a BGR format for display

    This is abstracted out to allow for experimentation

    Args:
        video: A numpy array with 1 byte per pixel, 3 channels RGB

    Returns:
        A numpy array with with 1 byte per pixel, 3 channels BGR
    """
    return video[:, :, ::-1]  # RGB -> BGR

def calculate_homography(self):
        # corners of depth
        point_depth=[[ 436 , 15 ], [ 101 , 287 ], [442, 475], [622,259], [44, 325], [119,363], [129, 344], [187, 101], [188, 257], [410, 254],[406, 102]]
        pts_depth = np.array(point_depth)

        # corners of image
        points_image=[[ 393 , 55 ], [ 98 , 293 ], [401, 476], [571, 272], [43, 339], [120,364],[122, 352], [175, 132],[184, 271],[383, 264],[379, 126]]
        pts_image = np.array(points_image)

        # Calculate Homography
        h, status = cv2.findHomography(pts_image,pts_depth)

        return h

def get_depth():
    return pretty_depth_cv(freenect.sync_get_depth()[0])

def get_video():
    return video_cv(freenect.sync_get_video()[0])

def get_contours(grey):
    contours, hierarchy = cv2.findContours(grey,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)[-2:]
    roi=[]
    for cnt in contours:
        x,y,w,h = cv2.boundingRect(cnt)
        g=grey.copy()
        roi.append(g[y:y+h,x:x+w])

        cv2.rectangle(grey,(x,y),(x+w,y+h),(200,0,0),2)

    return grey, roi

def select_slices(depth,start=1,slices_selected=4,slices=5):
    minimum=np.min(depth[depth>0])
    depth_masked=depth.copy()
    depth_masked[depth==255]=0
    maximum=np.max(depth_masked)
    step=(maximum-minimum)/slices
    threshold_1=(step*(start-1))+minimum
    threshold_2=threshold_1+(step*slices_selected)
    depth_masked[depth < threshold_1] = 0
    #depth_masked[depth >= threshold_1] = 255
    depth_masked[depth > threshold_2] = 0
    return depth_masked

def crop_image(image, zoom):
    image_h, image_w = image.shape
    h_desired = int(image_h*zoom/100)
    w_desired = int(image_w*zoom/100)

    zoomed=cv2.resize(image, (h_desired ,w_desired), interpolation=cv2.INTER_AREA)
    cv2.imshow("pippo", zoomed)
    cv2.waitKey(0)

    delta_h=int(h_desired-image_h/2)
    delta_w=int(w_desired-image_w/2)
    return  image[delta_h:image_h+delta_h, delta_w:image_w+delta_w]

def padding(image, border_size=10):
    return cv2.copyMakeBorder(image, border_size,border_size,border_size,border_size,cv2.BORDER_CONSTANT,value=0)

def crop(image, border_size=10):
    return image[border_size:image.shape[0]-border_size, border_size:image.shape[1]-border_size]

def get_point(image,border_size=10,tolerance=10):
    return (image.shape[0]-border_size-tolerance,image.shape[1]-border_size-tolerance)

def removing_bottom_floor(image):
    # Clear rows with uniform values (First floor removal)
    for i in range(image.shape[0]-1,-1, -1):
        row = image[i, :]
        if (np.max(row) - np.min(row)) < 8:
            image[i, :]=0
    return image

def get_neighbours(point, max_dimensions):
    l=[(-1, -1), (0, -1), (1, -1), (1, 0), (1, 1), (0, 1), (-1, 1), (-1, 0)]
    return {(point[0]+x,point[1]+y) for x,y in l if 0<=point[0]+x<max_dimensions[0] and 0<=point[1]+y<max_dimensions[1]}

def removing_floor_region_growing(image, output, border_size=10, point=None, tolerance=25):
    if not point:
        point = get_point(image,border_size)
    result=output.copy()
    value_ref=int(image[point[0],point[1]] )
    neighbours={point}
    visited=set()
    while len(neighbours)>0:
        p=neighbours.pop()
        n=set()
        n={x for x in get_neighbours(p, image.shape) if x not in visited and abs(int(image[x[0],x[1]])-value_ref)<tolerance}
        neighbours.update(n)
        result[p[0],p[1]]=0
        visited.add(p)
    return result

def removing_floor_connected_components(original, output, border_size=10, point=None, tolerance=8):
    if not point:
        point = get_point(original,border_size)
    image=original.copy()
    result=output.copy()
    value_ref=int(image[point[0],point[1]] )
    image[image<value_ref-tolerance]=0
    image[image>value_ref+tolerance]=0
    image[image>0]=1
    _, labels = cv2.connectedComponents(image)
    label=labels[point[0],point[1]]
    result[labels==label]=0
    return result

def live_mode():
    while True:
        # Original depth
        dep_original=get_depth()

        # Remove background
        depth_no_background=select_slices(dep_original)

        # Filtering depth
        depth_no_background_median = cv2.medianBlur(depth_no_background, 5)

        # Cropping image
        border_size=10
        depth_no_background_median_cropped = crop(depth_no_background_median,border_size=border_size)
        # Padding image for restore dimensions
        depth_no_background_median_cropped = padding(depth_no_background_median_cropped,border_size=border_size)

        # Sobel
        depth_sobel = cv2.Sobel(dep_original,cv2.CV_64F,0,1,ksize=5)
        abs_sobel64f = np.absolute(depth_sobel)
        depth_sobel = np.uint8(abs_sobel64f)

        # Blurring Sobel image
        kernel_size=9
        depth_sobel_blurred = cv2.blur(depth_sobel, (kernel_size,kernel_size))

        depth_without_floor=removing_floor_connected_components(depth_sobel_blurred,depth_no_background_median_cropped,border_size)
        depth_without_floor_blurred=cv2.medianBlur(depth_without_floor, 15)
        #depth_without_floor_blurred[depth_without_floor_blurred>0]=255
        cv2.imshow('Depth Without Backgroud and Floor', depth_without_floor_blurred)
        cv2.imshow('Depth Sobel', depth_sobel_blurred)

        if cv2.waitKey(1)==27:
            break

def batch_mode():
    from_kinect=True
    video=True
    path=os.path.dirname(__file__)
    path='/home/davide/LamIra/Test/segmentation_test'
    i=1
    # Defining starting element
    element = cv2.getStructuringElement(cv2.MORPH_CROSS,(3,3))
    depth_file=path+"/test/depth_"+str(i) #insert right number in i

    # Take depth image
    if from_kinect:
        dep_original=get_depth()
        v=get_video()
        cv2.imwrite(path+"/test/rgb_obj_4.png", v)
        cv2.imwrite(path+"/test/depth_obj_4.png", dep_original)
        exit(0)
    else:
        dep_original=cv2.imread(depth_file+"_starting.png", 0)
        v=None
        video=False

    # Remove background by mask
    depth_masked=select_slices(dep_original)

    # Filtering depth
    d_median = cv2.medianBlur(depth_masked, 5)

    border_size=10
    # Cropping image
    d_median_cropped = crop(d_median,border_size=border_size)
    # Padding image for restore dimensions
    d_median_cropped = padding(d_median_cropped,border_size=border_size)

    #d_edge = cv2.Canny(depth_masked,55,65,apertureSize = 5)
    #d_edge_2 = cv2.morphologyEx(d_edge, cv2.MORPH_GRADIENT, element)
    #d_edge_3 = cv2.medianBlur(d_edge_2, 5)
    #d_median_cropped = crop_image(depth_masked, 50)

    # Filtering Video
    if video:
        v_grey = cv2.cvtColor(v, cv2.COLOR_BGR2GRAY)
        v_edge=cv2.Canny(v_grey,70,90, apertureSize = 3)
        rgb_file=path+"/test/rgb_"+str(i) #insert right number in i

    # Save interediet file
    d_first_attempt = d_median_cropped.copy()
    d_final2 = select_slices(d_first_attempt, slices_selected=15, slices=16)

    depth_sobel = cv2.Sobel(dep_original,cv2.CV_64F,0,1,ksize=5)
    abs_sobel64f = np.absolute(depth_sobel)
    depth_sobel = np.uint8(abs_sobel64f)

    kernel_size=9
    depth_test = cv2.blur(depth_sobel, (kernel_size,kernel_size))
    #cv2.imshow("sobel", depth_sobel)
    #cv2.imshow("pippo", depth_test)

    image_without_floor=removing_floor_region_growing(depth_test,d_median_cropped,border_size)
    image_without_floor=cv2.medianBlur(image_without_floor, 15)
    #cv2.imshow("not floor", image_without_floor)
    d_median_cropped=image_without_floor

    # Clear columns with floor
    for k in range(d_median_cropped.shape[1]-1,-1, -1):
        column=d_median_cropped[:, k]
        obj=False
        for j,value in enumerate(column):
            if j == len(column)-1:
                break
            succ = column[j+1]
            if value > 0 and value < succ:
                obj=True
                break

        if not obj:
            d_median_cropped[:, k]=0

    d_median_cropped = cv2.medianBlur(d_median_cropped, 5)
    #d_median_cropped = cv2.morphologyEx(d_median_cropped, cv2.MORPH_GRADIENT, element)

    # Extracting ROI
    d_median_cropped, rois = get_contours(d_median_cropped)
    for l, roi in enumerate(rois):
        ret2,th2 = cv2.threshold(roi,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
        cv2.imwrite(depth_file+"_roi_" + str(l) + "_"+str(i)+".png", th2)

    # Saving Files
    cv2.imwrite(rgb_file+".png", v)
    cv2.imwrite(depth_file+".png", dep_original)
    '''
    cv2.imwrite(depth_file+"_without_background.png", depth_masked)
    cv2.imwrite(depth_file+"_final.png", d_first_attempt)
    cv2.imwrite(depth_file+"_final_2nd_select_slices.png", d_final2)
    cv2.imwrite(depth_file+"_floor_removal_by_column.png", d_median_cropped)
    '''

if __name__=="__main__":
    #live_mode()
    batch_mode()

max_stars_repo_name: marcolamartina/LamIra

max_issues_repo_name: marcolamartina/LamIra

max_forks_repo_name: marcolamartina/LamIra


text:
from scipy import ndimage as ndi

from skimage import io, feature
from skimage.color import rgb2gray
from skimage.feature import blob_dog
from scipy import ndimage
from copy import deepcopy

def img_open(filename):
        im = io.imread(filename)
        return im

def edge(filename):
        im = img_open(filename)

        # Zeroing red channel, combined blue + green channels into
        # single grayscale channel.
        im[:, :, 0] = 0
        im_grey = ndimage.gaussian_filter(rgb2gray(im), 2)

        # Compute the Canny filter for two values of sigma

        edges = feature.canny(im_grey, sigma=3)

        mask = ndi.binary_fill_holes(edges)
        return mask

def rm_back(mask, filename):
        im = img_open(filename)
        no_back = deepcopy(im)
        for x in range(mask.shape[0]):
                for y in range(mask.shape[1]):
                        if (mask[x, y] == 0):
                                no_back[x, y] = im[x, y] * mask[x, y]
        return no_back

def in_range(x, y, mask):
        # 40% hard coded
        centerx = mask.shape[0]/2
        centery = mask.shape[1]/2
        realx = abs(x-centerx)
        realy = abs(y-centery)
        if(realx**2 + realy**2 < (0.5*centerx)**2):
                return True
        else:
                return False

def blobs(im_no_back):
        im_no_back = rgb2gray(im_no_back)
        blob_list = blob_dog(im_no_back, min_sigma=2, threshold=0.6)
        blob_list[:, 2] = blob_list[:, 2] * 2**0.5
        return [blob_list]

max_stars_repo_name: mjip/producesort

max_issues_repo_name: mjip/producesort

max_forks_repo_name: mjip/producesort

reason: from scipy


text:
import csv
import os
import random
import cv2
from scipy import ndimage
import nibabel as nib
import numpy as np
import pandas as pd
import copy

def read_nifti_file(filepath):
    """Read and load volume"""
    # Read file
    scan = nib.load(filepath)
    # Get raw data
    scan = scan.get_fdata()
    return scan

def normalize(volume, bound=(-374, 426)):
    """Normalize the volume"""
    volume = np.clip(volume, bound[0], bound[1])
    volume = 255*((volume - bound[0]) / (bound[1] - bound[0]))
    volume = volume.astype("float32")
    return volume

def resize_volume(img):
    """Resize across z-axis"""
    # Set the desired shape
    desired_width = 224
    desired_height = 224

    # Get current depth
    current_width = img.shape[0]
    current_height = img.shape[1]

    # Compute depth factor
    width = current_width / desired_width
    height = current_height / desired_height

    width_factor = 1 / width
    height_factor = 1 / height
    # Rotate
    img = ndimage.rotate(img, 90, reshape=False)
    # Resize across z-axis
    img = ndimage.zoom(img, (width_factor, height_factor, 1), order=1)
    return img

def process_scan(path):
    """Read and resize volume"""
    # Read scan
    volume = read_nifti_file(path)
    # Normalize
    volume = normalize(volume)
    # Resize width, height and depth
    volume = resize_volume(volume)
    return volume

cur_csv_file = '/home/alec/renalData.csv'
new_csv_file = '/home/alec/renalData_new.csv'
path_dir = '/mnt/cadlabnas/datasets/RenalDonors'

data = pd.read_csv(cur_csv_file)
CLASSES = ['PC', 'SYNTHCON']
os.makedirs(f'{path_dir}/{CLASSES[0]}/img', exist_ok=True)
os.makedirs(f'{path_dir}/{CLASSES[1]}/img', exist_ok=True)

with open(new_csv_file, 'w') as f:
  writer = csv.DictWriter(f, fieldnames=['PC', 'PC_num', 'SYNTHCON', 'SYNTHCON_num', 'subset'])
  writer.writeheader()
  for row in data.iterrows():
    pc_img = row[1][CLASSES[0]]
    syn_img = row[1][CLASSES[1]]
    row_data = {}
    row_data['subset'] = row[1]['subset']
    for folder, filename in zip(CLASSES, [pc_img, syn_img]):
      if filename.endswith('.nii.gz'):
        full_path_filename = f'{path_dir}/{folder}/{filename}'
        scan = process_scan(full_path_filename)
        n_images = scan.shape[0]
        row_data[folder + '_num'] = n_images
        row_data[folder] = filename

        # for a nifti file, dim_0 is image width, dim_1 is height, dim_2 is channels/slices
        im_width, im_height, nifti_channels = np.asarray(scan).shape
        im_channels = 3
        half_c = im_channels // 3
        for i in range(n_images):
          # check bounds of data
          # within
          if (i - half_c >= 0) & (i + half_c <= nifti_channels - half_c):
              slice_start = i - half_c
              slice_end = i + half_c + 1
              # get temporary sub_slice
              t_sub_slice = scan[:,:,slice_start:slice_end]

          # outside left edge
          elif i - half_c < 0:
              # get temporary sub_slice
              t_sub_slice = scan[:,:,0:im_channels]

          # outside right edge
          elif i + half_c > nifti_channels - 1:
              # get temporary sub_slice
              t_sub_slice = scan[:,:,nifti_channels-im_channels:nifti_channels]

          # current shape of subslice is [width, height, channels]
          # need to change shape to [height, width, channels] for saving to disk/training
          curr_sub_slice = np.zeros([im_height, im_width, im_channels], dtype = np.uint8)

          # histogram equalization (adaptive with CLAHE)
          # run histEq on the normalized_9x_percentile midSlice
          ksize = 8
          clahe = cv2.createCLAHE(clipLimit=1.0, tileGridSize=(ksize, ksize))
          for channel_idx in range(im_channels):
              # transpose individual slice  ----- [width, height] --> [height, width]
              curr_slice = np.transpose(t_sub_slice[:,:,channel_idx])

              # normalize the slice
              # norm_9x_img = min_max_normalization(curr_slice, rquantile = 1)
              histEq_img = clahe.apply(np.asarray(curr_slice).astype(np.uint8))

              # place in position
              curr_sub_slice[:,:,channel_idx] = copy.deepcopy(histEq_img)

          # opencv saves images in BGR format (not RGB -- so interchange first and last images in sub-slice)
          tempRed = copy.deepcopy(curr_sub_slice[:,:,0])
          curr_sub_slice[:,:,0] = copy.deepcopy(curr_sub_slice[:,:,2])
          curr_sub_slice[:,:,2] = copy.deepcopy(tempRed)
          cv2.imwrite(f'{path_dir}/{folder}/img/{filename[:filename.find(".")]}_{i}.png', curr_sub_slice)

    writer.writerow(row_data)

max_stars_repo_name: alecpeltekian/ImgClassification

max_issues_repo_name: alecpeltekian/ImgClassification

max_forks_repo_name: alecpeltekian/ImgClassification

reason: import numpy,from scipy


text:
import matplotlib.pyplot as plt
import numpy as np
import proplot as pplt

from functions.formatting import *
from functions.base import *
from functions.logging import write_log

write_log(__file__)

output = 'save'
output_dir = current_dir / 'figures' / 'Fig_1'
output_dir.mkdir(parents=True, exist_ok=True)
fname = 'Fig_1a_isotopes'

h1 = np.array([6.9, 9.1, 5.9, 3.0, 0.7, 0.3])
h1 / h1.max()

h2 = np.array([2.1, 7.2, 9.2, 8.8, 5.5, 2.9, 0.9, 0.2])
h2 / h2.max()

sclf = 2
fig_width = 25*sclf
fig_height = 25*sclf

h1 = np.array([0.75824176, 1.        , 0.64835165, 0.32967033, 0.07692308, 0.03296703])
h2 = np.array([0.22826087, 0.7826087 , 1.        , 0.95652174, 0.59782609, 0.31521739,
               0.09782609, 0.02173913])

def lorentzian(x, x0, gamma):
    return np.divide(1, np.pi*gamma* (1 + ( (x-x0) / gamma  )**2))

x = np.linspace(0, 10, num=1000)

y1 = np.zeros_like(x)
for i, height in enumerate(h1):
    y1 += height*lorentzian(x, i+1, 0.01)

y2 = np.zeros_like(x)
for i, height in enumerate(h2):
    y2 += height*lorentzian(x, i+1, 0.01)

fig, axes = pplt.subplots(nrows=2, width=fig_width / 25.4, height=fig_height/25.4, wspace=0.01, hspace=0.05, sharey=False)
axes[0].plot(x, y1 / y1.max(), linewidth=1, color='k')
axes[0].format(ylabel='t$_{0}$')
axes[1].plot(x, y2 / y1.max(), linewidth=1, color='k')
axes[1].format(ylabel='t$_{1}$')
axes.format(xticks=[], yticks=[], xlabel='m/z')

if output == 'show':
    plt.show()
elif output == 'save':
    plt.savefig(output_dir / f'{fname}.png')
    plt.savefig(output_dir / f'{fname}.pdf')

max_stars_repo_name: Jhsmit/PyHDX-paper

max_issues_repo_name: Jhsmit/PyHDX-paper

max_forks_repo_name: Jhsmit/PyHDX-paper


text:
import os
import argparse
import re
from collections import namedtuple
import numpy as np
import torch
import random
import math

from latent_rationale.sst.constants import UNK_TOKEN, PAD_TOKEN
from latent_rationale.sst.plotting import plot_heatmap
from torch.nn.init import _calculate_fan_in_and_fan_out
from torch import nn

def get_device():
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')

def find_ckpt_in_directory(path):
    for f in os.listdir(os.path.join(path, "")):
        if f.startswith('model'):
            return os.path.join(path, f)
    print("Could not find ckpt in {}".format(path))

def filereader(path):
    """read SST lines"""
    with open(path, mode="r", encoding="utf-8") as f:
        for line in f:
            yield line.strip().replace("\\", "")

def tokens_from_treestring(s):
    """extract the tokens from a sentiment tree"""
    return re.findall(r"\([0-9] ([^\(\)]+)\)", s)

def token_labels_from_treestring(s):
    """extract token labels from sentiment tree"""
    return list(map(int, re.findall(r"\(([0-9]) [^\(\)]", s)))

Example = namedtuple("Example", ["tokens", "label", "token_labels"])

def sst_reader(path, lower=False):
    """
    Reads in examples
    :param path:
    :param lower:
    :return:
    """
    for line in filereader(path):
        line = line.lower() if lower else line
        line = re.sub("\\\\", "", line)  # fix escape
        tokens = tokens_from_treestring(line)
        label = int(line[1])
        token_labels = token_labels_from_treestring(line)
        assert len(tokens) == len(token_labels), "mismatch tokens/labels"
        yield Example(tokens=tokens, label=label, token_labels=token_labels)

def print_parameters(model):
    """Prints model parameters"""
    total = 0
    for name, p in model.named_parameters():
        total += np.prod(p.shape)
        print("{:24s} {:12s} requires_grad={}".format(name, str(list(p.shape)),
                                                      p.requires_grad))
    print("\nTotal parameters: {}\n".format(total))

def load_glove(glove_path, vocab, glove_dim=300):
    """
    Load Glove embeddings and update vocab.
    :param glove_path:
    :param vocab:
    :param glove_dim:
    :return:
    """
    vectors = []
    w2i = {}
    i2w = []

    # Random embedding vector for unknown words
    vectors.append(np.random.uniform(
        -0.05, 0.05, glove_dim).astype(np.float32))
    w2i[UNK_TOKEN] = 0
    i2w.append(UNK_TOKEN)

    # Zero vector for padding
    vectors.append(np.zeros(glove_dim).astype(np.float32))
    w2i[PAD_TOKEN] = 1
    i2w.append(PAD_TOKEN)

    with open(glove_path, mode="r", encoding="utf-8") as f:
        for line in f:
            word, vec = line.split(u' ', 1)
            w2i[word] = len(vectors)
            i2w.append(word)
            vectors.append(np.array(vec.split(), dtype=np.float32))

    # fix brackets
    w2i[u'-LRB-'] = w2i.pop(u'(')
    w2i[u'-RRB-'] = w2i.pop(u')')

    i2w[w2i[u'-LRB-']] = u'-LRB-'
    i2w[w2i[u'-RRB-']] = u'-RRB-'

    vocab.w2i = w2i
    vocab.i2w = i2w

    return np.stack(vectors)

def get_minibatch(data, batch_size=25, shuffle=False):
    """Return minibatches, optional shuffling"""

    if shuffle:
        print("Shuffling training data")
        random.shuffle(data)  # shuffle training data each epoch

    batch = []

    # yield minibatches
    for example in data:
        batch.append(example)

        if len(batch) == batch_size:
            yield batch
            batch = []

    # in case there is something left
    if len(batch) > 0:
        yield batch

def pad(tokens, length, pad_value=1):
    """add padding 1s to a sequence to that it has the desired length"""
    return tokens + [pad_value] * (length - len(tokens))

def prepare_minibatch(mb, vocab, device=None, sort=True):
    """
    Minibatch is a list of examples.
    This function converts words to IDs and returns
    torch tensors to be used as input/targets.
    """
    # batch_size = len(mb)
    reverse_map = None
    lengths = np.array([len(ex.tokens) for ex in mb])
    maxlen = lengths.max()

    # vocab returns 0 if the word is not there
    x = [pad([vocab.w2i.get(t, 0) for t in ex.tokens], maxlen) for ex in mb]
    y = [ex.label for ex in mb]

    x = np.array(x)
    y = np.array(y)

    if sort:  # required for LSTM
        sort_idx = np.argsort(lengths)[::-1]
        x = x[sort_idx]
        y = y[sort_idx]

        # to put back into the original order
        reverse_map = np.zeros(len(lengths), dtype=np.int32)
        for i, j in enumerate(sort_idx):
            reverse_map[j] = i

    x = torch.from_numpy(x).to(device)
    y = torch.from_numpy(y).to(device)

    return x, y, reverse_map

def plot_dataset(model, data, batch_size=100, device=None, save_path=".",
                 ext="pdf"):
    """Accuracy of a model on given data set (using minibatches)"""

    model.eval()  # disable dropout
    sent_id = 0

    for mb in get_minibatch(data, batch_size=batch_size, shuffle=False):
        x, targets, reverse_map = prepare_minibatch(
            mb, model.vocab, device=device, sort=True)

        with torch.no_grad():
            logits = model(x)

            alphas = model.alphas if hasattr(model, "alphas") else None
            z = model.z if hasattr(model, "z") else None

        # reverse sort
        alphas = alphas[reverse_map] if alphas is not None else None
        z = z.squeeze(1).squeeze(-1)  # make [B, T]
        z = z[reverse_map] if z is not None else None

        for i, ex in enumerate(mb):
            tokens = ex.tokens

            if alphas is not None:
                alpha = alphas[i][:len(tokens)]
                alpha = alpha[None, :]
                path = os.path.join(
                    save_path, "plot{:04d}.alphas.{}".format(sent_id, ext))
                plot_heatmap(alpha, column_labels=tokens, output_path=path)

            # print(tokens)
            # print(" ".join(["%4.2f" % x for x in alpha]))

            # z is [batch_size, num_samples, time]
            if z is not None:

                zi = z[i, :len(tokens)]
                zi = zi[None, :]
                path = os.path.join(
                    save_path, "plot{:04d}.z.{}".format(sent_id, ext))
                plot_heatmap(zi, column_labels=tokens, output_path=path)

            sent_id += 1

def xavier_uniform_n_(w, gain=1., n=4):
    """
    Xavier initializer for parameters that combine multiple matrices in one
    parameter for efficiency. This is e.g. used for GRU and LSTM parameters,
    where e.g. all gates are computed at the same time by 1 big matrix.
    :param w:
    :param gain:
    :param n:
    :return:
    """
    with torch.no_grad():
        fan_in, fan_out = _calculate_fan_in_and_fan_out(w)
        assert fan_out % n == 0, "fan_out should be divisible by n"
        fan_out = fan_out // n
        std = gain * math.sqrt(2.0 / (fan_in + fan_out))
        a = math.sqrt(3.0) * std
        nn.init.uniform_(w, -a, a)

def initialize_model_(model):
    """
    Model initialization.

    :param model:
    :return:
    """
    # Custom initialization
    print("Glorot init")
    for name, p in model.named_parameters():
        if name.startswith("embed") or "lagrange" in name:
            print("{:10s} {:20s} {}".format("unchanged", name, p.shape))
        elif "lstm" in name and len(p.shape) > 1:
            print("{:10s} {:20s} {}".format("xavier_n", name, p.shape))
            xavier_uniform_n_(p)
        elif len(p.shape) > 1:
            print("{:10s} {:20s} {}".format("xavier", name, p.shape))
            torch.nn.init.xavier_uniform_(p)
        elif "bias" in name:
            print("{:10s} {:20s} {}".format("zeros", name, p.shape))
            torch.nn.init.constant_(p, 0.)
        else:
            print("{:10s} {:20s} {}".format("unchanged", name, p.shape))

def get_predict_args():
    parser = argparse.ArgumentParser(description='SST prediction')
    parser.add_argument('--ckpt', type=str, default="path_to_checkpoint",
                        required=True)
    parser.add_argument('--plot', action="store_true", default=False)
    args = parser.parse_args()
    return args

def get_args():
    parser = argparse.ArgumentParser(description='SST')
    parser.add_argument('--save_path', type=str, default='sst_results/default')
    parser.add_argument('--resume_snapshot', type=str, default='')

    parser.add_argument('--num_iterations', type=int, default=-25)
    parser.add_argument('--batch_size', type=int, default=25)
    parser.add_argument('--eval_batch_size', type=int, default=25)
    parser.add_argument('--embed_size', type=int, default=300)
    parser.add_argument('--proj_size', type=int, default=300)
    parser.add_argument('--hidden_size', type=int, default=150)

    parser.add_argument('--lr', type=float, default=0.0002)
    parser.add_argument('--min_lr', type=float, default=1e-5)
    parser.add_argument('--lr_decay', type=float, default=0.5)
    parser.add_argument('--threshold', type=float, default=1e-4)
    parser.add_argument('--cooldown', type=int, default=5)
    parser.add_argument('--patience', type=int, default=5)
    parser.add_argument('--weight_decay', type=float, default=1e-5)
    parser.add_argument('--max_grad_norm', type=float, default=5.)

    parser.add_argument('--model',
                        choices=["baseline", "rl", "attention",
                                 "latent"],
                        default="baseline")
    parser.add_argument('--dist', choices=["", "hardkuma"],
                        default="")

    parser.add_argument('--print_every', type=int, default=100)
    parser.add_argument('--eval_every', type=int, default=-1)
    parser.add_argument('--save_every', type=int, default=-1)

    parser.add_argument('--dropout', type=float, default=0.5)

    parser.add_argument('--layer', choices=["lstm"], default="lstm")
    parser.add_argument('--train_embed', action='store_false', dest='fix_emb')

    parser.add_argument('--dependent-z', action='store_true',
                        help="make dependent decisions for z")

    # rationale settings for RL model
    parser.add_argument('--sparsity', type=float, default=0.0)
    parser.add_argument('--coherence', type=float, default=0.0)

    # rationale settings for HardKuma model
    parser.add_argument('--selection', type=float, default=1.,
                        help="Target text selection rate for Lagrange.")
    parser.add_argument('--lasso', type=float, default=0.0)

    # lagrange settings
    parser.add_argument('--lagrange_lr', type=float, default=0.01,
                        help="learning rate for lagrange")
    parser.add_argument('--lagrange_alpha', type=float, default=0.99,
                        help="alpha for computing the running average")
    parser.add_argument('--lambda_init', type=float, default=1e-4,
                        help="initial value for lambda")

    # misc
    parser.add_argument('--word_vectors', type=str,
                        default='data/sst/glove.840B.300d.sst.txt')
    args = parser.parse_args()
    return args