text:
import argparse
from mmdet.apis import init_detector, inference_detector
import mmcv
import cv2
import os
import glob
import numpy as np
#parser = argparse.ArgumentParser()
#parser.add_argument('--path', type=str, default='task3/dataset_path/')
#args = parser.parse_args()
def task4_main(path):
    # ##### person #####
    # config_file_person = "task4/configs/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco-person.py"
    # checkpoint_file_person = "task4/checkpoints/faster_rcnn_r50_fpn_1x_coco-person_20201216_175929-d022e227.pth"
    config_file_person = "configs/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco-person.py"
    checkpoint_file_person = "checkpoints/faster_rcnn_r50_fpn_1x_coco-person_20201216_175929-d022e227.pth"
    # ##################
    # ##### triage #####
    # config_file_triage = "task4/configs/triage_config/triage_config.py"
    # checkpoint_file_triage = "task4/work_dirs/0to3_500/ver7/epoch_41.pth"
    #config_file_triage = "configs/triage_config/triage_config.py"
    #checkpoint_file_triage = "work_dirs/triage1600/epoch_19.pth"
    config_file_triage = "configs/triage_config/triage_config_deformable_detr.py"
    checkpoint_file_triage = "work_dirs/deformableDETR/epoch_4.pth"
    # ##################
    # # build the model from a config file and a checkpoint file
    model_person = init_detector(config_file_person, checkpoint_file_person, device="cuda:0")
    model_triage = init_detector(config_file_triage, checkpoint_file_triage, device="cuda:0")
    # set(번호)_drone(번호)_triage(번호).jpg
    # set(1~5)_drone(1~3)_triage(1~3).jpg
    # im_folder = "dataset_path/set0"
    im_folder = path
    # 제출 레이블 형식 : set_num(1~5), drone_num(1~3), frame_name[사망, 긴급, 응급, 비응급]
    #import pdb; pdb.set_trace()
    person_results = []
    set_keys = ["set_1", "set_2", "set_3", "set_4", "set_5"]
    task4_answer = dict.fromkeys(set_keys)
    # #### Extract Region of Person ####
    for set_n in range(1,6):
        set_dict = dict()
        set_name = "set_"+ str(set_n)
        set_dir= im_folder + "set_0" + str(set_n) + "/"
        # print("set_dir : ",set_dir)

        for filename in glob.glob(set_dir): # filename : dataset_path/set_01/
            # print("filename: ", filename)
            drone_1, drone_2, drone_3 = dict(), dict(), dict()
            file_list = os.listdir(filename)
            file_list = sorted(file_list)
            triage_list=[]
            # set_dict_list = []
            for x in file_list:
                if 'triage' in x :
                    #file_list = os.listdir(filename) # ["set02_drone03_triage02.jpg", "set02_drone01_triage01.jpg", "set02_drone01_triage03.jpg", "set02_drone03_triage01.jpg", "set02_drone01_triage02.jpg", "set02_drone02_triage01.jpg", "set02_drone02_triage02.jpg"]
                    triage_list.append(x)
            drone_dict_list = dict()
            for file_idx in range(len(triage_list)):

                # print("file_idx: ", file_idx)
                answer_sheet = [0 for i in range(4)]
                ori_file = set_dir + triage_list[file_idx]
                # # print(file_list[i]) # set01_drone01_triage01.jpg
                # # print(set_dir + file_list[i]) # dataset_path/set01/set01_drone01_triage01.jpg
                result_person = inference_detector(model_person, ori_file) # person
                result_triage = inference_detector(model_triage, ori_file) # person
                labels = [
                    np.full(bbox.shape[0], idx, dtype=np.int32)
                    for idx, bbox in enumerate(result_triage)
                ]
                labels = np.concatenate(labels)
                bboxes_person = np.vstack(result_person)
                bboxes_triage = np.vstack(result_triage)
                # class-based NMS
                tag_is_in_person = [False for t in range(bboxes_triage.shape[0])]
                for p in range(bboxes_person.shape[0]):
                    person_pos = bboxes_person[p][:4] # each person has one or no triage tag.
                    is_exist_pos = []
                    max_score = -1
                    if bboxes_person[p][-1] < 0.5: continue # person post-processing threshold
                    for q in range(bboxes_triage.shape[0]):
                        triage_pos = bboxes_triage[q][:4]
                        if(triage_pos[0] >= person_pos[0]-50 and triage_pos[1] >= person_pos[1]-50 and triage_pos[2] <= person_pos[2]+50 and triage_pos[3] <= person_pos[3]+50):
                            is_exist_pos.append(q)
                            tag_is_in_person[q] = True
                            if bboxes_triage[q][-1] > max_score:
                                max_score = bboxes_triage[q][-1]
                    for k in range(len(is_exist_pos)):
                        if bboxes_triage[is_exist_pos[k]][-1] < max_score:
                            bboxes_triage[is_exist_pos[k]][-1] = 0
                for t in range(bboxes_triage.shape[0]):
                    if tag_is_in_person[t] == False:
                        bboxes_triage[t][-1] = 0
                scores = bboxes_triage[:, -1]
                score_thr = 0.3
                inds = scores > score_thr
                labels = labels[inds]
                for i in range(len(labels)):
                    answer_sheet[labels[i]] += 1
                drone_num = int(triage_list[file_idx].split("drone")[1][:2])
                drone = "drone_" + str(drone_num)
                img_key = triage_list[file_idx].split(".jpg")[0]
                if drone_num == 1:
                    drone_1[img_key] = answer_sheet
                elif drone_num == 2:
                    drone_2[img_key] = answer_sheet
                elif drone_num == 3:
                    drone_3[img_key] = answer_sheet

            drone_1_list, drone_2_list, drone_3_list = [], [], []
            drone_1_list.append(drone_1)
            drone_2_list.append(drone_2)
            drone_3_list.append(drone_3)

            set_drone1, set_drone2, set_drone3 = dict(), dict(), dict()
            set_drone1["drone_1"] = drone_1_list
            set_drone2["drone_2"] = drone_2_list
            set_drone3["drone_3"] = drone_3_list

            set_drone_list = []
            set_drone_list.append(set_drone1)
            set_drone_list.append(set_drone2)
            set_drone_list.append(set_drone3)
            task4_answer[set_name] = set_drone_list
            # print("\n \ntask4_answer[set_name] :\n ", task4_answer[set_name])
    print(task4_answer)
    #final_answer = dict()
    #final_answer["task4_answer"] = task4_answer
    return task4_answer
if __name__ == '__main__':
    path = 'dataset_path/'
    task4_main(path)

max_stars_repo_name: sojinleeme/mmdetection

max_issues_repo_name: sojinleeme/mmdetection

max_forks_repo_name: sojinleeme/mmdetection


text:
import torch
import random
import tiledspmm

from numpy import random
random.seed(123)
torch.manual_seed(123)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(123)

n = 1024
n_vals = n * n

def gen_coords(num_vals, dim1, dim2):
    coords = set()
    while len(coords) < num_vals:
        x, y = 0, 0
        while (x, y) == (0, 0):
            x, y = random.randint(0, dim1-1), random.randint(0, dim2-1)

        coords.add((x, y))
    return coords

def sparsify(coords, dim1, dim2):
    # return a sparse randomly generated tensor
    mat = torch.zeros(dim1, dim2, device=torch.device('cuda'))
    for x in range(0, dim1):
        for y in range(0, dim2):
            if (x, y) in coords:
                mat[x][y] = random.random()
    return mat

shapes = [
          ((n, n), (n, n)),
          ((n, 2*n), (2*n, n)),
          ((n, n), (n, int(n/2))),
          ((2*n, n), (n, int(n/2))),
        ]

# use these for sparse-sparse matrix multiplication
# b_coords = gen_coords(n_vals, b_rows, b_cols)
# b = sparsify(b_coords, n)

for i, ((a_rows, a_cols), (b_rows, b_cols)) in enumerate(shapes):
    # sparse dense matrix multiplication
    a_n_vals = (a_rows * a_cols) / n_vals # scales number of values to be (1/n_vals)% sparsity
    a_coords = gen_coords(a_n_vals, a_rows, a_cols)
    a = sparsify(a_coords, a_rows, a_cols)

    b = torch.rand((b_rows, b_cols), device=torch.device('cuda'))
    c = torch.zeros((a_rows, b_cols), device=torch.device('cuda'))

    exp = a@b

    a = a.to_sparse()
    vals = torch.Tensor.values(a).to('cpu')
    indices = torch.Tensor.indices(a).type(torch.IntTensor)
    rows = indices[0].to('cpu')
    cols = indices[1].to('cpu')
    nnz = len(vals)

    layer = "layer" + str(i)
    tiledspmm.tiledspmm_inspect_coo(a_rows, a_cols, b_cols, nnz, rows, cols, vals, layer)

    '''
    a = a.to_sparse_csr()
    vals = torch.Tensor.values(a).to('cpu')
    displ = torch.Tensor.crow_indices(a).type(torch.LongTensor).to('cpu')
    indices = torch.Tensor.col_indices(a).type(torch.LongTensor).to('cpu')

    tiledspmm.tiledspmm_inspect_csr(a_rows, a_cols, b_cols, displ, indices, vals)
    '''
    tiledspmm.tiledspmm_mm(b, c, layer)

    if torch.allclose(exp, c):
        print(f'\nTest {i} passed!\n')
    else:
        print(f'\nTest {i} failed!\n')
        print('expected: ', exp)
        print('ours(sparse): ', c)
        print('# nonzero in exp:', torch.nonzero(exp).shape)
        print('# nonzero in ours:', torch.nonzero(c).shape)

tiledspmm.tiledspmm_clean()

ext: py

max_stars_repo_name: smoorjani/matrix-multiplication

max_issues_repo_name: smoorjani/matrix-multiplication

max_forks_repo_name: smoorjani/matrix-multiplication

reason: from numpy


text:
import numpy as np
import pandas as pd
import re
import warnings
from typing import Iterable

# Feature Engineering functions
def create_lag(df, features, lag=[1], droplagna=True):
    """create additional lagged features

    Args:
        df (pd.DataFrame): the input data.
        features (Iterable, str): the features in ``df.columns`` that you you want to create lags
        lag (list, optional): lags that you want to create. Defaults to [1]. For example, if lag=[1,3,9], then create feature of each features with lag 1,3, and 9.
        droplagna (bool, optional): drop the missing records/rows due to lags. Notics that this doesn't drop all na records but only those due to shift/lags. Defaults to True.

    Returns:
        pd.DataFrame: new dataframe contains lags features as columns. This doesn't contain original df content.
    """
    # settup
    if isinstance(features, str):
        features = [features]
    if isinstance(lag, int):
        lag = [lag]

    # make the shift
    resList = []
    for i in lag:
        res = df[features].shift(i)
        suffix = 'lag' if i > 0 else 'lead'
        i = i if i > 0 else -i
        res.columns = [str(s)+'_{}{}'.format(suffix, i)  for s in features ]
        resList.append(res)
    res = pd.concat(resList, axis=1)

    if droplagna:
        lagmax, lagmin = max(lag), min(lag)
        if lagmax > 0 and lagmin > 0:
            res = res.iloc[lagmax:, :]
            print('#lags={} is dropped'.format(lagmax))
        elif lagmax > 0 and lagmin < 0:
            res = res.iloc[lagmax:, :]
            res = res.iloc[:lagmin, :]
            print('#lags={} and {} is dropped'.format(lagmax, lagmin))
        elif lagmax < 0 and lagmin < 0:
            res = res.iloc[:lagmin, :]
            print('#lags={} is dropped'.format(lagmin))
        else:
            raise ValueError('Impossible case where lagmin > lagmax')

    return res

def create_lead(df, features: Iterable, lead=[1], dropleadna=True):
    """create additional leads features. This function is eqivalent to ``create_lag(...,lag=[-l for l in lead],...)``

    Args:
        df (pd.DataFrame): the input data.
        features (Iterable, str): the features in ``df.columns`` that you you want to create lags
        lead (list, optional): leads that you want to create. Defaults to [1]. For example, if lead=[1,3,9], then create feature of each features with lead 1,3, and 9.
        droplagna (bool, optional): drop the missing records/rows due to lags. Notics that this doesn't drop all na records but only those due to shift/lags. Defaults to True.

    Returns:
        pd.DataFrame: new dataframe contains leads features as columns. This doesn't contain original df content.
    """
    lag = -np.array(lead)
    return create_lag(df, features, lag=lag, droplagna=dropleadna)

def add_MA(df, features, window, **kwargs):

    if isinstance(features, str):
        features = [features]
    if isinstance(window, int):
        window = [window]

    resList = []
    for w in window:
        res = df[features].rolling(w, min_periods=1, **kwargs).mean()
        res.columns = [str(s)+f'_ma{w}' for s in features]
        resList.append(res)
    res = pd.concat(resList, axis=1)

    return pd.concat([df, res], axis=1), res.columns

def add_EMA(df, features, alpha, adjust=True, **kwargs):

    if isinstance(features, str):
        features = [features]
    if isinstance(alpha, int):
        alpha = [alpha]

    resList = []
    for a in alpha:
        res = df[features].ewm(alpha=a, min_periods=1,
                             adjust=adjust, **kwargs).mean()
        res.columns = [str(s)+'_ema{:g}'.format(a*10) for s in features]
        resList.append(res)
    res = pd.concat(resList, axis=1)
    return pd.concat([df, res], axis=1), res.columns

def add_MS(df, features, window, **kwargs):

    if isinstance(features, str):
        features = [features]
    if isinstance(window, int):
        window = [window]

    resList = []
    for w in window:
        res = df[features].rolling(w, min_periods=1, **kwargs).std()
        res.columns = [str(s)+f'_ms{w}' for s in features]
        # use feature mean to fill na  (shall be the first 3 records)
        res.fillna(res.mean(), inplace=True)
        resList.append(res)
    res = pd.concat(resList, axis=1)

    return pd.concat([df, res], axis=1), res.columns

def add_skew(df, features, window, **kwargs):
    """
    Args:
        df ([type]): [description]
        features ([type]): [description]
        window ([type]): should be at least 3. Otherwise returned are all NaN.

    Returns:
        [type]: [description]
    """
    if isinstance(features, str):
        features = [features]
    if isinstance(window, int):
        window = [window]

    resList = []
    for w in window:
        res = df[features].rolling(w, min_periods=1, **kwargs).skew()
        res.columns = [str(s)+f'_skew{w}' for s in features]
        # use feature mean to fill na (shall be the first 3 records)
        res.fillna(res.mean(), inplace=True)
        resList.append(res)
    res = pd.concat(resList, axis=1)
    return pd.concat([df, res], axis=1), res.columns

def npshift(arr, num, fill_value=np.nan):
    result = np.empty_like(arr)
    if num > 0:
        result[:num] = fill_value
        result[num:] = arr[:-num]
    elif num < 0:
        result[num:] = fill_value
        result[:num] = arr[-num:]
    else:
        result[:] = arr
    return result

def percent_change(data):
    """
    Calculate the %change between the last value and the mean of previous values.
    Makes data more comparable if the abs values of the data change a lot.
    Usage: df.rolling(window=20).aggregate(percent_change)
    """
    previous_values = data[:-1]
    last_value = data[-1]
    return (last_value - np.mean(previous_values)) / np.mean(previous_values)

def fourier_features(index, freq, order):
    """
    `Reference`_: https://www.kaggle.com/ryanholbrook/seasonality
    Example:
    >>> # Compute Fourier features to the 4th order (8 new features) for a
    >>> # series y with daily observations and annual seasonality:
    >>> # fourier_features(y, freq=365.25, order=4)
    """
    time = np.arange(len(index), dtype=np.float32)
    k = 2 * np.pi * (1 / freq) * time
    features = {}
    for i in range(1, order + 1):
        features.update({
            f"sin_{freq}_{i}": np.sin(i * k),
            f"cos_{freq}_{i}": np.cos(i * k),
        })
    return pd.DataFrame(features, index=index)

def replace_outliers(df, window=7, k=3, method='mean'):
    """
    Inputs:
        df: pd.DataFrame/Series
        window: rolling window size to decide rolling mean/std/median.
        k: multiplier of rolling std.
    Return:
        pd.DataFrame of the same shape as input df.
    """
    # de-mean
    mean_series = df.rolling(window).mean()[window-1:]
    abs_meandiff = (df[window-1:] - mean_series).abs()
    std_series = df.rolling(window).std()[window-1:]
    median_series = df.rolling(window).median()[window-1:]
    # identify >k(=3 in default) standard deviations from zero
    this_mask = abs_meandiff > (std_series * k)
    tmp = df[:window-1].astype(bool)
    tmp.values[:] = False
    this_mask = pd.concat([tmp, this_mask], axis=0)
    # Replace these values with the median accross the data
    if method == 'median':
        to_use = median_series
    elif method == 'mean':
        to_use = mean_series
    else:
        raise ValueError(f'method {method} not found.')
    return df.mask( this_mask, to_use )

def recover_from_diff(prediction_diff, y, dlist):
    """
    Transform difference-predicted series back to original
    prediction level via y.
    Only support first difference case.

    Parameters
    ----------
    prediction_diff: model.predict() result of training or testing data
    y: the original level of the feature that is the target of
       the prediction_diff.
       Usually, y should have same length as prediction_diff.
    dlist: list that stores 1st-difference record for each column of y.
       currently, the approach only deals with d=0,1 in dlist.

    Returns
    -------
    prediction: recovered prediction series with tail entry truncated
       since the final prediction is out of bound. And usually the first
       prediction correspond to y[1] instead of y[0].
    """
    import numpy as np

    assert (np.array(dlist) <= 1).all(), "Currently, the method doesn't\
        deal with difference case higher than 1 (in dlist)."
    yarray = np.array(y, ndmin=1)
    parray = np.array(prediction_diff, ndmin=1)
    assert len(dlist) == yarray.shape[-1]

    res = parray.copy()
    for i in range(len(dlist)):
        # pass if no first difference on the current column
        if dlist[i] == 0:
            continue
        # concept: y_hat[i+1] = y[i] + delta_hat(y[i+1]-y[i])
        res[:, i] = yarray[:, i] + parray[:, i]
    return res.reshape(y.shape)

# # need testing this function, the advanced version of the original.
# def recover_from_diff(prediction_diff, y, dlist, step=1):
#      """
#     Transform difference-predicted series back to original
#     prediction level via y.
#     TODO: this function only supports recovery where the prediction is based on ground truth, ie,
#     the function doesn't support the case where the prediction is generated through the exogenous
#     feature prediction procedure

#     Parameters
#     ----------
#     prediction_diff: model.predict() result of training or testing data
#     y:  the original level of the feature that is the target of
#         the prediction_diff.
#         Usually, y should have same length as prediction_diff.
#     dlist: list - stores 1st-difference record for each column of y.
#         currently, the approach only deals with d=0,1 in dlist.
#     step: int - how many step ahead is the prediction compared to y.

#     Eg, if pred is 2-step-ahead on 1st-diff level, then step=2, dlist=[1]
#     Eg, if pred is 1-step-ahead on 2st-diff level, then step=1, dlist=[2]

#     Returns
#     -------
#     prediction: recovered prediction series with tail entry truncated
#        since the final prediction is out of bound. And usually the first
#        prediction correspond to y[1] instead of y[0].
#     """
#     import numpy as np

#     yarray = np.array(y, ndmin=1)
#     parray = np.array(prediction_diff, ndmin=1)
#     assert len(dlist) == yarray.shape[-1]

#     res = parray.copy()
#     # loop over each column
#     for i in range(len(dlist)):
#         # continue if no diff() at all on the current column
#         if dlist[i] == 0:
#             continue
#         # loop for recovering prediction back to 0th-hierarchy level
#         for hierarchy in range(dlist[i], 0, -1):
#             # get corresponding y_true series for 1 step diff() back transformation
#             helper = np.diff(yarray[:, i], n=hierarchy-1)
#             # locate the position: not sure and need double check
#             gap = helper.shape[0] - res.shape[0]
#             res[:, i] = helper[-res.shape[0]-gap-step:-gap-step] + res[:, i]

#     return res.reshape(y.shape)

#%% DataMelter
class DataMelter():

    def __init__(self):
        pass

    @staticmethod
    def _concat_back(frame: pd.DataFrame, frame2: pd.DataFrame, ):
        assert frame2.shape[0] % frame.shape[0] == 0, "frame is not divisible by frame2 wrt nrows."
        multiplier = int(frame2.shape[0] / frame.shape[0])
        augmented_frame = pd.concat([frame]*multiplier, axis=0)
        assert (augmented_frame.index == frame2.index).all(), "index of the two frames are not the same."
        return pd.concat([augmented_frame, frame2], axis=1)

    @staticmethod
    def melt_features(frame: pd.DataFrame, id_cols=None, input_cols=None,
                  input2dummy_value={}, dummy_name: str=None, value_name: str=None, ignore_index=False, concat_back=False):
        """
        Unpivot a dataframe wide to long format. This function is an extension of ``pd.melt`` that provides additionally the mapper for input column names to dummy values.

        :param frame: the input dataframe.
        :type frame: pd.DataFrame
        :param id_cols: column(s) to use as identifier variables, defaults to None
        :type id_cols: [tuple, list, or ndarray], optional
        :param input_cols: [Columns to unpivot. If not specified, used all columns that are not set as ``id_vars``], defaults to None
        :type input_cols: [tuple, list, or ndarray], optional
        :param input2dummy_value: [Mapper that convert ``input_cols`` names to dummy values under the column ``dummy_name``], defaults to {}
        :type input2dummy_value: dict, optional
        :param dummy_name: [Name to use for the `variable` column. If None it uses ``frame.columns.name`` or `variable`.], defaults to None
        :type dummy_name: str, optional
        :param value_name: [Name to use for the ‘value’ column.], defaults to None
        :type value_name: str, optional
        :param ignore_index: [whether to ignore index when melting the data], defaults to False.
        :type ignore_index: bool, optional
        :param concat_back: [whether to concat the melted features back to original dataframe], defaults to False.
        :type concat_back: bool, optional
        :return: [Unpivoted DataFrame.]
        :rtype: [pd.DataFrame]
        """

        check = pd.Index(input_cols).isin(frame.columns)
        assert check.all(), f"Cannot find columns {input_cols[~check]} in frame."
        # prepare
        value_vars = input2dummy_value.values()
        if len(value_vars) == 0:
            value_vars = input_cols
        # output
        res = frame.rename(
            columns=input2dummy_value
        ).melt(
            id_vars=id_cols, value_vars=value_vars,
            var_name=dummy_name, value_name=value_name,
            ignore_index=ignore_index
        )

        # concat back if True
        if concat_back:
            res = DataMelter._concat_back(frame.drop(columns=input_cols), res, )

        return res

    @staticmethod
    def melt_features_regex(frame: pd.DataFrame, features: Iterable[str],
        dummy_value: Iterable[str]=None, dummy_name: Iterable[str]=None,
        value_name: Iterable[str]=None, ignore_index=False):
        """
        This method is similar to ``melt_features`` but can melt different group of features via regex.

        :param frame: [the input dataframe that already contains the lead features]
        :type frame: pd.DataFrame
        :param features: [the target featues regex list for prediction task]
        :type features: Iterable[str].

        :param targets_suffix_regex: [the suffix name regex pattern for targets_stem]
        :type targets_suffix_regex: str, optional

        :param dummy_name:
        :type dummy_name: Iterable[str], defaults to None.

        :param value_name:
        :type value_name: Iterable[str], defaults to None.

        :param ignore_index:
        :type ignore_index: bool, defaults to False.

        :return: [description]
        :rtype: [type]
        """
        # initial check
        assert not isinstance(features, str), "targets_stem should be Iterable, not str."

        if isinstance(dummy_name, str) or dummy_name is None:
            dummy_name = [dummy_name] * len(features)
        else:
            assert len(dummy_name) == len(features)

        if dummy_value is None:
            dummy_value = features * len(features)
        elif isinstance(dummy_value, str):
            dummy_value = [dummy_value] * len(features)
        else:
            assert len(dummy_value) == len(features)

        if isinstance(value_name, str):
            value_name = [value_name] * len(features)
        elif value_name is None:
            value_name = features
        else:
            assert len(value_name) == len(features)

        # melt targets down to long format
        collector = []
        for i, feature_pattern in enumerate(features):
            cols = frame.filter(regex=feature_pattern).columns
            tomap = cols.str.extract(f"({dummy_value[i]})", expand=False)
            long_df = DataMelter.melt_features(
                frame, input_cols=cols,
                input2dummy_value=dict(zip(cols, tomap)),
                dummy_name=dummy_name[i], value_name=value_name[i]
            )
            collector.append(long_df)

        # prepare output
        output = pd.concat(collector, axis=1)
        uniq_cols = output.columns.drop_duplicates()
        return output , uniq_cols

# %%

max_stars_repo_name: wkCircle/myPythonLibrary

max_issues_repo_name: wkCircle/PythonToolKit

max_forks_repo_name: wkCircle/PythonToolKit

reason: import numpy


text:
'''
@Author:        ZM
@Date and Time: 2020/7/26 14:37
@File:          mnist_dataset.py
'''

import numpy as np
import gzip

def get_mnist(root_dir='D:/datasets/MNIST'):
    train_images_path = f'{root_dir}/train-images-idx3-ubyte.gz'
    train_labels_path = f'{root_dir}/train-labels-idx1-ubyte.gz'

    val_images_path = f'{root_dir}/t10k-images-idx3-ubyte.gz'
    val_labels_path = f'{root_dir}/t10k-labels-idx1-ubyte.gz'

    with gzip.open(train_images_path, 'rb') as f:
        X_train = np.frombuffer(f.read(), dtype='uint8', offset=16).reshape(-1, 28, 28)

    with gzip.open(train_labels_path, 'rb') as f:
        Y_train = np.frombuffer(f.read(), dtype='uint8', offset=8)

    with gzip.open(val_images_path, 'rb') as f:
        X_val = np.frombuffer(f.read(), dtype='uint8', offset=16).reshape(-1, 28, 28)

    with gzip.open(val_labels_path, 'rb') as f:
        Y_val = np.frombuffer(f.read(), dtype='uint8', offset=8)

    return (X_train, Y_train), (X_val, Y_val)


text:
from collections import defaultdict, namedtuple
from itertools import combinations
from multiprocessing import Pool
from multiprocessing.managers import SharedMemoryManager
from functools import partial

import numpy as np
from scipy.sparse import lil_matrix, csr_matrix

from SuperPang.lib.utils import read_fasta, write_fasta, reverse_complement, print_time
from SuperPang.lib.Compressor import Compressor

import graph_tool as gt
from graph_tool.all import Graph

from mappy import Aligner

### see if we can put GS.clear_filters() inside the set filters method, instead of having a lot of them throughout the code

class Assembler:

    ### MULTIPROCESS METHODS
    # In order to save memory minimize serialization/deserialization overheads when multiprocessing we will pass the input variables to our target function as class attributes (a.k.a. gloryfied globals).
    # We set them before opening the process pool: this way they get inherited when forking the original process, and linux's Copy On Write makes it so memory is actually not copied unless it is modified.
    # The output of our function still needs to be serialized by the workers and deserialized by the main process, so this will not work if outputs are large and frequent
    # For that we can use shared memory (see __init__), particularly for output amenable to be stored in array (ints, bytes...)

    multiprocessing_globals = tuple()

    @classmethod
    def set_multiprocessing_globals(cls, *args):
        # Multiprocessing.Pool.map comes with a significant overhead in serializing the arguments and sending them to the workers, which beats the purpose of multiprocessing.
        # What we'll do instead is store the required variables as a class attribute before starting the pool. Those will be copied when forking the process, which is much faster.
        cls.multiprocessing_globals = args

    @classmethod
    def clear_multiprocessing_globals(cls):
        cls.multiprocessing_globals = tuple()

    @classmethod
    def get_multiprocessing_globals(cls):
        return cls.multiprocessing_globals if len(cls.multiprocessing_globals) > 1 else cls.multiprocessing_globals[0]

    @classmethod
    def multimap(cls, fun, threads, iterable, *args, single_thread_threshold = 24, imap = False, pool = None):
        args = [iterable] + list(args)
        cls.set_multiprocessing_globals(*args)
        if threads == 1 or len(iterable) < single_thread_threshold:
            res = list(map(fun, range(len(iterable))))
        else:
            with Pool(threads) as pool:
                res = pool.map(fun, range(len(iterable)))
        cls.clear_multiprocessing_globals()
        return res
    ###############################

    ##    @profile
    def __init__(self, fasta, ksize, threads):
        """
        Build a De-Bruijn Graph from an input fasta file
        """

        self.ksize = ksize
        self.includedSeqs = 0
        self.hash2vertex = {}
        self.vertex2hash = {}
        self.seqPaths = {} # name: pathSet

        ### Read input sequences
        print_time(f'Reading sequences')
        self.seqDict = read_fasta(fasta, Ns = 'split')
        nNs = len({name.split('_Nsplit_')[0] for name in self.seqDict})

        ### Kmer hashing related parameters
        self.compressor = Compressor(self.ksize)
        clength = self.ksize // 4 + self.ksize % 4                # the length of a hash
        maxseqlength = max(len(s) for s in self.seqDict.values()) # length of the largest input sequence
        maxnkmers = maxseqlength-self.ksize+1                     # number of kmers in the largest input sequence

        ### Get kmers from sequences and create DBG
        print_time(f'Creating DBG from {len(self.seqDict)} sequences ({nNs} originals)')
        edges = set()
        totalSize = sum(len(seq) for seq in self.seqDict.values())
        elapsedSize = 0
        lastSize = 0

        with SharedMemoryManager() as smm:
            # seqMem and rcSeqMem are two shared mem buffers that will hold the fwd and rev sequences
            #    they are pre-allocated using the size of the largest sequence, for each sequence we will
            #    store/recover the correct amount of bytes according to its length
            # hashMem and revhashMem are two continuous shared mem buffers that will hold all the hashes concatenated
            #    they are pre-allocated using the number of kmers in the largest sequence multiplied by the hash size
            #    to store/recover a particular hash we can just use the right slice of the mem buffer, since hash size is constant

            seqMem     = smm.SharedMemory(size=maxseqlength)
            rcSeqMem   = smm.SharedMemory(size=maxseqlength)
            hashMem    = smm.SharedMemory(size=clength*maxnkmers)
            revhashMem = smm.SharedMemory(size=clength*maxnkmers)
            self.set_multiprocessing_globals(seqMem, rcSeqMem, hashMem, revhashMem, self.ksize, clength, self.compressor)

            with Pool(threads) as pool:

                for name, seq in self.seqDict.items():
                    elapsedSize += len(seq)
                    if len(seq) <= self.ksize:
                        continue
                    else:
                        self.includedSeqs += 1
                        rcSeq = reverse_complement(seq)
                        seq, rcSeq = seq.encode(), rcSeq.encode()
                        nkmers = len(seq)-self.ksize+1

                        seqMem.buf  [:len(seq)  ] = seq
                        rcSeqMem.buf[:len(rcSeq)] = rcSeq

                        seq2hashes_ = partial(self.seq2hashes, seqlength=len(seq)) # we can't pass seqlength as a global when forking
                                                                                #  since it changes for every sequence
                        pool.map(seq2hashes_, range(nkmers))

                        # Retrieve the kmer hashes from the memory buffers
                        # Note the use of a generator expression so that all the hashes are not residing in memory at the same time
                        #  ... so we can't close the pool until we finish going through such generator expression!
                        hashes = (hashMem.buf[i*clength:(i+1)*clength].tobytes() for i in range(nkmers))
                        revhashes = (revhashMem.buf[i*clength:(i+1)*clength].tobytes() for i in range(nkmers))

                        self.clear_multiprocessing_globals()

                        # Populate the DBG
                        sp = set()
                        prev = -1 # store the previous vertex here so we can build (prev, v) edge tuples
                        for h, rh in zip(hashes, revhashes):
                            if h not in self.hash2vertex and rh not in self.hash2vertex:
                                v = len(self.hash2vertex)
                                self.hash2vertex[h] = v
                                self.vertex2hash[v] = h
                            elif h in self.hash2vertex:
                                v = self.hash2vertex[h]
                            elif rh in self.hash2vertex:
                                v = self.hash2vertex[rh]
                            else:
                                raise('wtf')
                            sp.add(v)
                            if prev >=0:
                                edges.add( (prev, v) )
                                edges.add( (v, prev) )
                            prev = v

                        self.seqPaths[name] = sp

                        if elapsedSize - lastSize > totalSize/100:
                            print_time(f'\t{round(100*elapsedSize/totalSize, 2)}% bases added, {len(self.hash2vertex)} vertices, {len(edges)} edges         ', end = '\r')
                            lastSize = elapsedSize

        self.G = Graph()
        self.G.add_edge_list(edges)
        print_time(f'\t100% bases added, {self.G.num_vertices()} vertices, {self.G.num_edges()} edges         ')
        n2m = {n: f'seq{i}' for i, n in enumerate(self.seqDict)}
        self.m2n = {m: n for n, m in n2m.items()}
        #write_fasta({n2m[n]: self.seqDict[n] for n in self.seqDict}, fasta + '.parsed.fasta')
        #self.ref = Aligner(fn_idx_in = fasta + '.parsed.fasta', n_threads = 3)

##    @profile
    def run(self, minlen, mincov, bubble_identity_threshold, genome_assignment_threshold, threads):
        """
        Create contigs from a De-Bruijn Graph
        """

        Contig = namedtuple('Contig', ['scaffold', 'i', 'seq', 'tseq', 'cov', 'origins', 'successors'])
        contigs = {}
        addedPaths = set()
        addedSeqs = set()

        ### Identify connected components
        print_time('Identifying connected components')
        vertex2comp = {v: c for v, c in enumerate(gt.topology.label_components(self.G, directed = False)[0])}
        comp2vertices = defaultdict(set)
        currentScaffold = 0
        for v, c in vertex2comp.items():
            comp2vertices[c].add(v)

        ### Process connected components
        for c, vs in comp2vertices.items():

            print_time(f'Working on comp {c+1}/{len(comp2vertices)}')

            addedSeqs_thiscomponent = set() # store the sequences overlapping with our paths here
                                            #  when we finish processing the DBG component we will
                                            #  use this to update addedSeqs
            # Get the seqPaths that were not overlapping with previous DBG components
            #  assuming that the same sequence can't overlap in two different DBG graph components
            seqPathsRemaining = {name: vs for name, vs in self.seqPaths.items() if name not in addedSeqs}

            ### Reconstruct non-branching paths
            print_time('\tCollecting non-branching paths')
            paths = []
            path_limits = defaultdict(list)
            inits = set()
            for v in vs:
                if len(set(self.G.get_all_neighbors(v)) - {v}) != 2: # -{v} to avoid self loops being init points
                    inits.add(v)

            # Break cycle if required
            isCycle = False
            ignoreEdge = set()
            if not inits:
                #is this a cycle?
                if min([len(set(self.G.get_all_neighbors(v))) for v in vs]) == 2: # no sources or sinks
                    for v in vs:
                        neighbors = self.G.get_all_neighbors(v)
                        if len(set(neighbors)) == 2: # find the first extender edge, and break there
                            inits.add(v)
                            inits.add(neighbors[0])
                            isCycle = True
                            ignoreEdge = {v, neighbors[0]}
                            break

            for j, ini in enumerate(inits):
                for n in set(self.G.get_all_neighbors(ini)):
                    p = [ini, n]
                    last = ini
                    while True:
                        s = set(self.G.get_all_neighbors(p[-1])) - {last, p[-1]} #p[-1] to avoid stopping at self loops
                        if len(s) != 1:
                            break
                        last = p[-1]
                        s = s.pop()
                        if isCycle: # Break cycle at the pre-defined point
                            if last in ignoreEdge and s in ignoreEdge:
                                break
                        p.append(s)
                    p = tuple(p)
                    paths.append(p)

            # Corner case
            if len(paths) == 2: #### case in which only a starting/ending vertex is common to two sequences,
                                #### so both are joined by that vertex but that vertex seems an extender (1 in 1 out)
                                #### If there was really one path we are cutting it, but we'll join it later
                p = paths[0] # cut it in half
                if len(p)%2:
                    cut = int((len(p)+1)/2)
                    inits.add(p[cut-1])
                    paths = [p[:cut], p[:cut][::-1], p[-cut:], p[-cut:][::-1]]

            # Check that there are no len==1 paths
            for p in paths:
                assert len(p) > 1

            ### Build a graph of connected non-branching paths, and split into fwd and rev component
            print_time(f'\tBuilding sequence graph out of {len(paths)} paths')

            # Remove paths with terminally duplicated vertices
            # Those broke our code once. we had three vSs with paths (0,0,1,2) (0,1,2) and (2,1,0) under the same pset (see below)
            paths = [p for p in paths if p[0] != p[1] and p[-1] != p[-2]]

            # Translate paths into sequences
            path2seq = dict(zip(paths, self.multimap(self.reconstruct_sequence, threads, paths, self.vertex2hash, self.compressor)))

            # Collect prefixes-suffixes for locating potential overlaps
            start2paths = defaultdict(set) # first k nucleotides
            sStart = {}
            sEnd = {}
            for path, seq in path2seq.items():
                start, end = seq[:self.ksize], seq[-self.ksize:]
                start2paths[start].add(path)
                sStart[path] = start
                sEnd[path] = end

            # Build sequence graph
            GS = Graph()
            path2vertexGS = {}
            good = 0
            lessGood = 0

            for p in paths:
                assert p not in path2vertexGS
                path2vertexGS[p] = int(GS.add_vertex())
            for p1 in paths:
                for p2 in start2paths[sEnd[p1]]:
                    GS.add_edge(path2vertexGS[p1], path2vertexGS[p2])

            vertex2path = {v: p for p,v in path2vertexGS.items()}

            ### Extract forward component
            print_time('\tExtracting forward component')
            comp2vertexGS = defaultdict(set) ## ahora uso v para todos los vértices. debería usar v para vértices de bruijn, y otra terminología para vertices secuencia. YA HE CAMBIADO PARTE
            comps =  gt.topology.label_components(GS, directed = False)[0]
            GS.vp.comps = comps
            GS.vp.vfilt = GS.new_vertex_property('bool')
            GS.ep.efilt = GS.new_edge_property('bool')

            for vS in GS.vertices():
                comp2vertexGS[GS.vp.comps[vS]].add(int(vS))

            # Try to find separated components in the seq graph that contain the same vertices in the kmer graph (thus are RC of each other)
            rcComps = {}
            for cS1, cS2 in combinations(comp2vertexGS,2):
                v1 = {v for vS in comp2vertexGS[cS1] for v in vertex2path[vS]}
                v2 = {v for vS in comp2vertexGS[cS2] for v in vertex2path[vS]}
                if v1 == v2:
                    rcComps[cS1] = cS2
                    rcComps[cS2] = cS1

            # If we didn't find a RC for all components this means that in some case fwd and rev are joined in the same component
            noRC = {cS for cS in comp2vertexGS if cS not in rcComps}
            for cS in noRC: # for each of those evil components...
                psets = defaultdict(list)
                for vS in comp2vertexGS[cS]:
                    pset = frozenset(vertex2path[vS])
                    psets[pset].append(vS)
                for vSs in psets.values():
                    assert len(vSs) == 2

                assert len(psets) == len(comp2vertexGS[cS]) / 2 # each sequence path should have a reverse complement equivalent (same vertices in the kmer graph, reverse order)
                vS2rc = {}
                for vS1, vS2 in psets.values():
                    vS2rc[vS1] = vS2
                    vS2rc[vS2] = vS1
                self.set_vertex_filter(GS, comp2vertexGS[cS])
                assert {int(vS) for vS in GS.vertices()} == comp2vertexGS[cS] # did we alter vS indices when filtering?
                assert(len(set(gt.topology.label_components(GS, directed = False)[0]))) == 1 # check that this is only one component

                for pset, (vS1, vS2) in psets.items():
                    assert vertex2path[vS1] == vertex2path[vS2][::-1]

                # Map the input sequences to nodes in the sequence graph (in the original orientation of the input sequence)
                names = list(seqPathsRemaining.keys())
                seqPathsGS = dict(zip(names, self.multimap(self.get_seqPathsGS, 24, names, self.seqPaths, psets, path2seq, vertex2path, self.seqDict)))
                seqPathsGS = {name: spGSs for name, spGSs in seqPathsGS.items() if spGSs}

                addedNodes = set()
                for spGS in seqPathsGS.values():
                    for vS in spGS:
                        addedNodes.update( (vS, vS2rc[vS]) )

                # Add nodes in the sequence graph that weren't fully contained in an original sequence
                for vS in comp2vertexGS[cS] - addedNodes:
                    fakeName = hash( (len(seqPathsGS), vS) )
                    seqPathsGS[fakeName] = {vS} | set(GS.get_all_neighbors(vS)) # add its neighbors so that it'll overlap with the real sequences in at least one vS

                seqPathsGS_rev = {name: {vS2rc[vS] for vS in vSs} for name, vSs in seqPathsGS.items()}

                x = {vS for vSs in seqPathsGS.values() for vS in vSs}
                y = {vS2rc[vS] for vS in x}

                # Separate fwd and rev
                vSs1 = set()
                refSeqNames = sorted(seqPathsGS, key = lambda name: seqPathsGS[name]) # sort by decreasing number of nodes in the GS

                addedNames = {refSeqNames[0]}
                vSs1.update(seqPathsGS[refSeqNames[0]]) # start from the longest seq

                while True:
                    remaining = [n for n in refSeqNames if n not in addedNames] # keep them sorted
                    if not remaining:
                        break
                    best_name = ''
                    best_ol = -1 # -1 so that even if there are no overlaps (there is more than one subcomponent, and the current one is full) we add something
                    for name in remaining: # keep starting from the longest seq, so that if there are no overlaps we add the first longest one
                        fwd = seqPathsGS[name]
                        rev = seqPathsGS_rev[name]
                        fwd_ol = len(fwd & vSs1) / len(fwd) # find the best overlap (looking in both orientations) to the existing fwd component
                        rev_ol = len(rev & vSs1) / len(rev)
                        if fwd_ol > best_ol and fwd_ol >= rev_ol:
                            best_ol, best_name, best_vSs = fwd_ol, name, fwd
                        elif rev_ol > best_ol and rev_ol > fwd_ol:
                            best_ol, best_name, best_vSs = rev_ol, name, rev
                    assert best_name
                    addedNames.add(best_name)
                    best_vSs = {vS for vS in best_vSs if vS2rc[vS] not in vSs1} # don't add reverse complements at this stage
                                                                                # some are really there and needed to keep everything into a single component
                                                                                # some will be duplications (e.g. because one of the input genomes has an inverted region)
                    vSs1.update(best_vSs)

                vSs2 = {vS2rc[vS] for vS in vSs1}
                v1 = {v for vS in vSs1 for v in vertex2path[vS]}
                v2 = {v for vS in vSs2 for v in vertex2path[vS]}

                assert (vSs1 | vSs2) == comp2vertexGS[cS]

                # Identify subcomponents in vS1
                GS.clear_filters()
                self.set_vertex_filter(GS, vSs1)
                subcomps1 = defaultdict(set)
                for vS,c in zip(GS.vertices(), gt.topology.label_components(GS, directed = False)[0]):
                    subcomps1[c].add(vS)
                GS.clear_filters()
                # Identify subcomponents in vS2
                self.set_vertex_filter(GS, vSs2)
                subcomps2 = defaultdict(set)
                for vS,c in zip(GS.vertices(), gt.topology.label_components(GS, directed = False)[0]):
                    subcomps2[c].add(vS)
                GS.clear_filters()
                # Identify RC components
                assert len(subcomps1) == len(subcomps2)

                subcomp2vertex1 = {scS1: {v for vS in svS1 for v in vertex2path[vS]} for scS1, svS1 in subcomps1.items()} # so we don't have to compute this inside the loop
                subcomp2vertex2 = {scS2: {v for vS in svS2 for v in vertex2path[vS]} for scS2, svS2 in subcomps2.items()}

                # Try to merge into one fwd and one rev subcomp
                # In some cases the fwd and rev sequences can be actually be present in the same genome
                # So we identify the largest fwd subcomponent, throw in the rest of components in fwd and rev, and hope that they become connected when everything is together
                largestSubcomp = list(sorted(subcomps1, key = lambda scS: len(subcomps1[scS]), reverse = True))[0]
                combined = subcomps1[largestSubcomp]
                for scS in subcomps1:
                    if scS != largestSubcomp:
                        fwd = subcomps1[scS]
                        rev = {vS2rc[vS] for vS in fwd}
                        combined = combined | fwd | rev

                subcomps1_corrected = defaultdict(set)
                subcomps2_corrected = defaultdict(set)
                self.set_vertex_filter(GS, combined)
                for vS, c  in zip(GS.vertices(), gt.topology.label_components(GS, directed = False)[0]):
                    subcomps1_corrected[c].add(vS)
                    subcomps2_corrected[c].add(vS2rc[vS])

                # The lines above may make a rev subcomponent may end up being contained in a fwd component
                # E.g. let us have initially three components A, B, C (fwd) and A', B', C' (rev)
                # We correct and it turns out that B' got inside A so by the code above we end up with AB', B, C (fwd) and A'B, B', C' (rev)
                # So we want to remove B from fwd and B' from rev, so that we don't end up having duplicate paths
                bad1 = set()
                bad2 = set()
                for scS1, svS1 in subcomps1_corrected.items():
                    for scS2, svS2 in subcomps2_corrected.items():
                        if   svS2.issubset(svS1) and not svS1.issubset(svS2):
                            bad2.add(scS2)
                        elif svS1.issubset(svS2) and not svS2.issubset(svS1):
                            bad1.add(scS1)
                subcomps1_corrected = {c: vSs for c, vSs in subcomps1_corrected.items() if c not in bad1}
                subcomps2_corrected = {c: vSs for c, vSs in subcomps2_corrected.items() if c not in bad2}

                # Keep going
                subcomps1 = subcomps1_corrected
                subcomps2 = subcomps2_corrected

                subcomp2vertex1 = {scS1: {v for vS in svS1 for v in vertex2path[vS]} for scS1, svS1 in subcomps1.items()} # so we don't have to compute this inside the loop
                subcomp2vertex2 = {scS2: {v for vS in svS2 for v in vertex2path[vS]} for scS2, svS2 in subcomps2.items()}

                # Assertions
                v1 = set.union(*list(subcomp2vertex1.values()))
                v2 = set.union(*list(subcomp2vertex2.values()))
                assert v1 == v2

                for svS1, svS2 in combinations(subcomps1.values(), 2):
                    assert not svS1 & svS2 # if they are not RC they should not share any paths
                for svS1, svS2 in combinations(subcomps2.values(), 2):
                    assert not svS1 & svS2 # if they are not RC they should not share any paths

                for scS1, svS1 in subcomps1.items():
                    for cS2, vSs in comp2vertexGS.items(): # they should not share any paths with other components in the sequence graph either
                        if cS2 != cS:
                            assert not vSs & svS1
                for scS2, svS2 in subcomps1.items():
                    for cS2, vSs in comp2vertexGS.items(): # they should not share any paths with other components in the sequence graph either
                        if cS2 != cS:
                            assert not vSs & svS2

                GS.clear_filters()

                first = True
                for scS1, svS1 in subcomps1.items():
                    v1 = subcomp2vertex1[scS1]
                    for scS2, svS2 in subcomps2.items():
                        v2 = subcomp2vertex2[scS2]
                        if v1 == v2:
                            i1 = cS if first else len(comp2vertexGS) # recycle the previous component index for cS1
                            comp2vertexGS[i1] = svS1
                            first = False
                            i2 = len(comp2vertexGS)
                            comp2vertexGS[i2] = svS2
                            rcComps[i1] = i2
                            rcComps[i2] = i1
                            break
                    else: # assert that all our subcomponents have RC components
                        assert False

                GS.clear_filters()

            assert len(rcComps) == len(comp2vertexGS) # all our components have a rc component, splitting was successful

            # Treat the component with more ATG as fwd and report it
            ATGs = defaultdict(int)
            for cS, vSs in comp2vertexGS.items():
                for vS in vSs:
                    ATGs[cS] += path2seq[vertex2path[vS]].count('ATG')
            added = set()

            compS2paths = defaultdict(set)
            for cS1, cS2 in rcComps.items():
                if cS1 in added:
                    continue
                fwdComp = cS1 if ATGs[cS1] >= ATGs[cS2] else cS2
                fwdPaths = [vertex2path[vS] for vS in comp2vertexGS[fwdComp]]
                for p in fwdPaths:
                    compS2paths[currentScaffold].add(p)

                added.add(cS1)
                added.add(cS2)
                currentScaffold += 1

            # Check that we included all the input vertices
            assert vs == {v for p in paths for v in p}
            # Check that no paths appear in two different scaffolds
            for cs1, cs2 in combinations(compS2paths, 2):
                assert not compS2paths[cs1] & compS2paths[cs2]

            ### Process each scaffold in the sequence graph
            for scaffold, paths in compS2paths.items():
                # Test that the scaffolds are really connected
                GS.clear_filters() # NEVER FORGET!!!
                self.set_vertex_filter(GS, {path2vertexGS[p] for p in paths})
                msg = f'\tScaffold {scaffold}, {GS.num_vertices()} seq vertices, {GS.num_edges()} seq edges'
                print_time(msg, end = '\r')
                assert len(set(gt.topology.label_components(GS, directed = False)[0])) == 1

                # Compute scaffold length and skip if too short
                scaffoldLen = 0
                for p in paths:
                    seq = path2seq[p]
                    if GS.get_out_neighbors(path2vertexGS[p]).shape[0]:
                        seq = seq[:-self.ksize] # trim the overlap unless this is a terminal node
                    scaffoldLen += len(seq)
                msg += f', {scaffoldLen} bases'
                if minlen and scaffoldLen < minlen:
                    msg += ' (< minlen, IGNORED)'
                    print_time(msg)
                    continue

                if mincov:
                    # Compute scaffold coverage
                    scaffoldCov = np.mean(self.multimap(self.vertex2originslen, threads, [v for p in paths for v in p], seqPathsRemaining, single_thread_threshold = 10000))
                    msg += f', cov {round(scaffoldCov, 2)}'
                    if mincov and scaffoldCov < mincov:
                        msg += ' (< mincov, IGNORED)'
                        print_time(msg)
                        continue

                print_time(msg, end='\r')

                # Identify and flatten bubbles
                bubble_paths = set()
                bubble_vertex2origins = defaultdict(set) # Keep track of the new ones here, the originals are calculated on the fly to save memory
                if bubble_identity_threshold and bubble_identity_threshold < 1:
                    path_limits = defaultdict(list)
                    for p in paths:
                        path_limits[(p[0], p[-1])].append(p)
                    m1 = 0
                    for (start, end), ps in path_limits.items():
                        ps = sorted(ps, key = lambda x: len(x), reverse = True)
                        minpathlen = len(ps[0])
                        if len(ps) > 1: # more than one path with similar start and end
                            p1 = ps[0] # the longest one
                            s1 = path2seq[p1]
                            ref = Aligner(seq = s1, n_threads = 1)
                            newOris = set()
                            for p2 in ps[1:]:                       # note that I am considering the 1st and last kmer here, which will always be shared!
                                al = list(ref.map(path2seq[p2]))[0] # this results in higher identities for short bubbles, which is not the worst
                                if al.mlen/al.blen >= bubble_identity_threshold:
                                    bubble_paths.add(p2)
                                    m1 += 1
                                    newOris.update( {ori for v2 in p2 for ori in self.vertex2origins(v2, seqPathsRemaining)} )
                            for v1 in p1:
                                bubble_vertex2origins[v1].update(newOris) # count vertices in p1 as appearing in all of the origins in the bubble paths that we just discarded

                    paths = [p for p in paths if p not in bubble_paths]
                    msg += f', removed {m1} bubbles'
                    print_time(msg, end = '\r')
                    GS.clear_filters()
                    self.set_vertex_filter(GS, {path2vertexGS[p] for p in paths})

                # Join contiguous unbranching paths (can happen after solving the fwd component and popping bubbles
                predecessors = defaultdict(set)
                successors = defaultdict(set)
                for p1 in paths:
                    for vS2 in GS.get_out_neighbors(path2vertexGS[p1]):
                        if vertex2path[vS2] not in bubble_paths:
                            successors[p1].add(vertex2path[vS2])
                    for vS2 in GS.get_in_neighbors(path2vertexGS[p1]):
                        if vertex2path[vS2] not in bubble_paths:
                            predecessors[p1].add(vertex2path[vS2])

                extenders = {p for p in paths if len(successors[p]) <= 1 and (successors[p] or predecessors[p])}
                extenders = extenders | {p for p in paths if len(predecessors[p]) == 1} ### added later, consider removing this line if we start failing assertions
                joinStarters =  {p for p in extenders if len(predecessors[p]) != 1 or len(successors[list(predecessors[p])[0]]) > 1}

                if not joinStarters: # Is this a cycle
                    if min(len(ps) for ps in predecessors.values()) == 1: # This does not cover the case in which there is a cycle but also a linear component
                        joinStarters = {list(extenders)[0]}

                joinExtenders = {p for p in extenders if p not in joinStarters}
                assert extenders == joinStarters | joinExtenders

                paths = [p for p in paths if p not in joinStarters | joinExtenders] # this also removes some joinStarters that can't be extended (no successors), but we'll add them back in the next bit
                newPaths = []

                visited = set()
                for start in joinStarters:
                    new_path = start
                    end = start
                    succs = successors[start]
                    succ = list(succs)[0] if succs else None
                    extended = False
                    while succ in joinExtenders:
                        extended = True
                        visited.add(succ)
                        new_path += succ[1:]
                        end = succ
                        succs = successors[succ]
                        succ = list(succs)[0] if succs else None
                    paths.append(new_path)
                    if extended:
                        sStart[new_path] = sStart[start]
                        sEnd[new_path] = sEnd[end]
                        newPaths.append(new_path)
                        predecessors[new_path] = predecessors[start]
                        for pred in predecessors[new_path]:
                            successors[pred].add(new_path)
                        successors[new_path] = successors[end]
                        for succ in successors[new_path]:
                            predecessors[succ].add(new_path)

                path2seq.update(dict(zip(newPaths, self.multimap(self.reconstruct_sequence, threads, newPaths, self.vertex2hash, self.compressor))))

                assert visited == joinExtenders

                # Update scaffold graph
                pathSet = set(paths)
                for p, preds in predecessors.items():
                    predecessors[p] = [pred for pred in preds if pred in pathSet]
                predecessors = defaultdict(list, [(p, preds) for p, preds in predecessors.items() if p in pathSet])
                for p, succs in successors.items():
                    successors[p]   = [succ for succ in succs if succ in pathSet]
                successors   = defaultdict(list, [(p, succs) for p, succs in successors.items()   if p in pathSet])

                # Check that paths with one predecessor and successor happen only around branching points
                for p in paths:
                    if len(predecessors[p])==1 and len(successors[p])==1:
                        pred = predecessors[p][0]
                        succ = successors[p][0]
                        assert len(successors[pred]) > 1 or len(predecessors[pred]) > 1
                        assert len(successors[succ]) > 1 or len(predecessors[succ]) > 1

                msg += f', processing {len(paths)} paths'
                print_time(msg, end = '\r')

                # By default we'll trim overlaps at 3', but we make exceptions
                # The idea is to remove also overlaps between two sequences sharing the same predecessor
                trim_left  = set()
                keep_right = set()
                for p in paths:
                    if len(successors[p]) > 1:
                        for succ in successors[p]:
                            if len(succ) > 1:
                                trim_left.add(succ)
                        keep_right.add(p)
                for p in paths:
                    if not successors[p] or len(p) < 2 or (len(p)<=2 and p in trim_left):
                        keep_right.add(p)
                    for succ in successors[p]:
                        if succ in trim_left:
                            keep_right.add(p)

                trimmedPaths = []
                for p in paths:
                    tp = tuple(p) # copy
                    if p not in keep_right:
                        tp = tp[:-1]
                    if p in trim_left:
                        tp = tp[1:]
                    assert tp
                    trimmedPaths.append(tp)

                vs = {v for p in paths for v in p}
                vst = {v for p in trimmedPaths for v in p}
                assert vs == vst

                trimmedSeqs = []
                for p in paths:
                    tseq = path2seq[p]
                    if p not in keep_right:
                        tseq = tseq[:-self.ksize]
                    if p in trim_left:
                        tseq = tseq[self.ksize:]
                    trimmedSeqs.append(tseq)

                # Update global results

                path2id  = {p: (scaffold, i) for i, p in enumerate(paths)}
                ori2covs = self.multimap(self.get_ori2cov, threads, trimmedPaths, seqPathsRemaining, bubble_vertex2origins)
                # Using trimmedPaths/trimmedSeqs: don't take into account overlaps with other paths for mean cov calculation and origin reporting
                #ori2covs = self.multimap(self.get_ori2cov, threads, trimmedSeqs, self.m2n, self.ref)

                foundOris = set()

                for p, ori2cov, tseq in zip(paths, ori2covs, trimmedSeqs):
                    assert p not in addedPaths
                    id_ = path2id[p]
                    # Even for complete genomes, we can have non-branching paths that belong to the core but don't have full coverage in some sequences
                    # This can happen if there are missing bases at the beginning or end of some of the input sequences (contigs in the original assemblies)
                    # So there is no branching, but some are missing
                    avgcov = sum(ori2cov.values())
                    goodOris = {ori for ori, cov in ori2cov.items() if cov >= genome_assignment_threshold}
                    goodOris.add( list(sorted(ori2cov, key = lambda ori: ori2cov[ori], reverse = True))[0] ) # add at least the origin with the highest cov
                    goodOris = {ori.split('_Nsplit_')[0] for ori in goodOris}
                    # Go on
                    contigs[id_] = Contig(scaffold   = id_[0],
                                          i          = id_[1],
                                          seq        = path2seq[p],
                                          tseq       = tseq,
                                          cov        = avgcov,
                                          origins    = goodOris,
                                          successors = [path2id[succ] for succ in successors[p]])
                    addedPaths.add(p)
                    for ori in ori2cov:
                        addedSeqs_thiscomponent.add(ori)

                msg += ', done'
                print_time(msg)

            # Wait until we have finished processing a whole DBG component for this
            addedSeqs.update(addedSeqs_thiscomponent)

        ### Finished!
        return contigs

    ################################ AUX METHODS

    @staticmethod
    def set_vertex_filter(GS, vSS):
        """
        Set a vertex filter in a graph_tool Graph
        """
        for vS in GS.vertices():
            GS.vp.vfilt[vS] = vS in vSS
        GS.set_vertex_filter(GS.vp.vfilt)

    @classmethod
    def seq2hashes(cls, i, seqlength):
        """
        Obtain sequence hashes for the kmers of a forward and a reverse sequence
        Stores results in two shared memory buffers
        """
        # hashMem and revhashMem are two continuous shared mem buffers that will hold all the hashes concatenated
        # for each index, we just read/write from/to the proper slices, since hash size is constant
        seqMem, rcSeqMem, hashMem, revhashMem, ksize, clength, compressor = cls.get_multiprocessing_globals()
        seq, rcSeq = seqMem.buf[:seqlength], rcSeqMem.buf[:seqlength]
        ri = len(rcSeq) - i - ksize
        hashMem.buf[i*clength:(i+1)*clength] = compressor.compress(seq[i:i+ksize])
        revhashMem.buf[i*clength:(i+1)*clength] = compressor.compress(rcSeq[ri:ri+ksize])

    @classmethod
    def reconstruct_sequence(cls, i):
        """
        Reconstruct a nucleotide sequence for a De-Bruijn Graph path
        """
        paths, vertex2hash, compressor = cls.get_multiprocessing_globals()
        path = paths[i]
        kmers = []
        vertex2kmer = partial(cls.vertex2kmer, vertex2hash = vertex2hash, compressor = compressor)
        # Each vertex corresponds to two kmers (fwd and rev)
        # Which kmers give us a meaningful sequence?
        # We are reading this in a given direction (from path[0] to path[-1])
        # For the starting kmer (in position 0), pick the one that overlaps with any of the two kmers in position 1
        k0x = vertex2kmer(path[0])
        k1x = vertex2kmer(path[1])
        for k1 in (k1x, reverse_complement(k1x)):
            for k0 in (k0x, reverse_complement(k0x)):
                if k1[:-1] in k0:
                    kmers.append(k0)

        # Then just keep advancing in the path, for each vertex pick the kmer that overlaps with the previously-picked kmer
        for v in path[1:]:
            k = vertex2kmer(v)
            ks = (k, reverse_complement(k))
            toadd = []
            for k in ks:
                if k[:-1] == kmers[-1][1:]:
                    toadd.append(k)
            assert len(toadd) <= 2
            if not toadd:
                break
            kmers.extend(toadd)
        return cls.seq_from_kmers(kmers)

    @classmethod
    def vertex2kmer(cls, v, vertex2hash, compressor):
        """
        Return the kmer corresponding to a given vertex in the De-Bruijn Graph
        """
        return compressor.decompress(vertex2hash[v])

    @staticmethod
    def seq_from_kmers(kmers):
        """
        Reconstruct a nucleotide sequence from a list of kmers
        """
        fullseq = [kmers[0]]
        extra = [kmer[-1] for kmer in kmers[1:]]
        fullseq.extend(extra)
        return ''.join(fullseq)

    @classmethod
    def get_seqPathsGS(cls, i):
        """
        Return the vertices in the Sequence Graph that are contained in a given input sequence in its original orientation
        """
        names, seqPaths, psets, path2seq, vertex2path, seqDict = cls.get_multiprocessing_globals()
        name = names[i]
        path = seqPaths[name]

        spGS = set()
        for pset, (vS1, vS2) in psets.items():
            if pset.issubset(path):
                for vS in (vS1, vS2):
                    if path2seq[vertex2path[vS]] in seqDict[name]: # check the string to make sure that we only add nodes from the same orientation
                        spGS.add(vS)
                        break
        return spGS

    @classmethod
    def seqInVerticesiREMOVE(cls, i):
        """
        Return whether a sequence shares at least one vertex with a set of vertices
        """
        seqs, seqPaths, vertices = cls.get_multiprocessing_globals()
        seq = seqs[i]
        res = False
        for v in seqPaths[seq]:
            if v in vertices:
                res = True
            break
        return res

    @classmethod
    def get_ori2cov(cls, i):
        """
        For a given path in the De-Bruijn Graph, calculate the fraction of its vertices that are contained in each input sequence
        """
        paths, seqPaths, bubble_vertex2origins = cls.get_multiprocessing_globals()
        path = paths[i]
        ori2cov = defaultdict(int)
        for v in path:
            for ori in cls.vertex2origins(v, seqPaths) | bubble_vertex2origins[v]: # note how I'm also getting origins not really associated to this path
                ori2cov[ori] += 1                                                         # but to bubbles that were originally in this path before we popped them
        return {ori: cov/len(path) for ori, cov in ori2cov.items()}

    @classmethod
    def get_ori2covN(cls, i):
        """
        For a given sequence, calculate the fraction that is contained in each input sequence
        """
        seqs, m2n, ref = cls.get_multiprocessing_globals()
        seq = seqs[i]
        return {m2n[align.ctg]: align.mlen / len(seq) for align in ref.map(seq)}

    @classmethod
    def vertex2origins(cls, v, seqPaths):
        """
        Return the names of the input sequences that contain a given vertex in the De-Bruijn Graph
        """
        return {name for name in seqPaths if v in seqPaths[name]}

    @classmethod
    def vertex2originslen(cls, i):
        """
        Return the number of input sequences that contain a given vertex in the De-Bruijn Graph
        """
        vertices, seqPaths = cls.get_multiprocessing_globals()
        return len(cls.vertex2origins(vertices[i], seqPaths))

max_stars_repo_name: fpusan/SuperPang

max_issues_repo_name: fpusan/SuperPang

max_forks_repo_name: fpusan/SuperPang

reason: import numpy,from scipy


text:
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from utils.util import opt_get
from models.modules import Basic, thops

class HCFlowNet_SR(nn.Module):
    def __init__(self, opt, step=None):
        super(HCFlowNet_SR, self).__init__()
        self.opt = opt
        self.quant = opt_get(opt, ['quant'], 256)

        hr_size = opt_get(opt, ['datasets', 'train', 'GT_size'], 160)
        hr_channel = opt_get(opt, ['network_G', 'in_nc'], 3)
        scale = opt_get(opt, ['scale'])

        if scale == 4:
            from models.modules.FlowNet_SR_x4 import FlowNet
        elif scale == 8:
            from models.modules.FlowNet_SR_x8 import FlowNet
        else:
            raise NotImplementedError('Scale {} is not implemented'.format(scale))

        # hr->lr+z
        self.flow = FlowNet((hr_size, hr_size, hr_channel), opt=opt)

        self.quantization = Basic.Quantization()

    # hr: HR image, lr: LR image, z: latent variable, u: conditional variable
    def forward(self, hr=None, lr=None, z=None, u=None, eps_std=None,
                add_gt_noise=False, step=None, reverse=False, training=True):

        # hr->z
        if not reverse:
            return self.normal_flow_diracLR(hr, lr, u, step=step, training=training)
        # z->hr
        else:
            return self.reverse_flow_diracLR(lr, z, u, eps_std=eps_std, training=training)

    #########################################diracLR
    # hr->lr+z, diracLR
    def normal_flow_diracLR(self, hr, lr, u=None, step=None, training=True):
        # 1. quantitize HR
        pixels = thops.pixels(hr)

        # according to Glow and ours, it should be u~U(0,a) (0.06 better in practice), not u~U(-0.5,0.5) (though better in theory)
        hr = hr + (torch.rand(hr.shape, device=hr.device)) / self.quant
        logdet = torch.zeros_like(hr[:, 0, 0, 0]) + float(-np.log(self.quant) * pixels)

        # 2. hr->lr+z
        fake_lr_from_hr, logdet = self.flow(hr=hr, u=u, logdet=logdet, reverse=False, training=training)

        # note in rescaling, we use LR for LR loss before quantization
        fake_lr_from_hr = self.quantization(fake_lr_from_hr)

        # 3. loss, Gaussian with small variance to approximate Dirac delta function of LR.
        # for the second term, using small log-variance may lead to svd problem, for both exp and tanh version
        objective = logdet + Basic.GaussianDiag.logp(lr, -torch.ones_like(lr)*6, fake_lr_from_hr)

        nll = ((-objective) / float(np.log(2.) * pixels)).mean()

        return torch.clamp(fake_lr_from_hr, 0, 1), nll

    # lr+z->hr
    def reverse_flow_diracLR(self, lr, z, u, eps_std, training=True):

        # lr+z->hr
        fake_hr = self.flow(z=lr, u=u, eps_std=eps_std, reverse=True, training=training)

        return torch.clamp(fake_hr, 0, 1)

max_stars_repo_name: CJWBW/HCFlow

max_issues_repo_name: CJWBW/HCFlow

max_forks_repo_name: CJWBW/HCFlow


text:
import numpy as np
import mxnet as mx
from mxnet.gluon import nn, HybridBlock
from mxnet.util import use_np
from autogluon_contrib_nlp.utils.config import CfgNode
from autogluon_contrib_nlp.layers import get_activation, get_norm_layer
from .. import constants as _C

@use_np
class BasicMLP(HybridBlock):
    def __init__(self, in_units,
                 mid_units,
                 out_units,
                 num_layers=1,
                 normalization='layer_norm',
                 norm_eps=1E-5,
                 dropout=0.1,
                 data_dropout=False,
                 activation='leaky',
                 weight_initializer=None,
                 bias_initializer=None,
                 prefix=None, params=None):
        """
        data -> [dropout] * (0/1) -> [Dense -> Normalization -> ACT] * N -> dropout -> Dense -> out

        Parameters
        ----------
        in_units
        mid_units
        out_units
        num_layers
            Number of intermediate layers
        normalization
        norm_eps
        dropout
        activation
        """
        super().__init__(prefix=prefix, params=params)
        self.in_units = in_units
        self.data_dropout = data_dropout
        if mid_units < 0:
            mid_units = in_units
        with self.name_scope():
            self.proj = nn.HybridSequential()
            with self.proj.name_scope():
                if num_layers > 0 and data_dropout:
                    self.proj.add(nn.Dropout(dropout))
                for i in range(num_layers):
                    self.proj.add(nn.Dense(units=mid_units,
                                           in_units=in_units,
                                           flatten=False,
                                           weight_initializer=weight_initializer,
                                           bias_initializer=bias_initializer,
                                           use_bias=False))
                    self.proj.add(get_norm_layer(normalization,
                                                 axis=-1,
                                                 epsilon=norm_eps,
                                                 in_channels=mid_units))
                    self.proj.add(get_activation(activation))
                    in_units = mid_units
                self.proj.add(nn.Dropout(dropout))
                self.proj.add(nn.Dense(units=out_units,
                                       in_units=in_units,
                                       weight_initializer=weight_initializer,
                                       bias_initializer=bias_initializer,
                                       flatten=False))

    def hybrid_forward(self, F, x):
        return self.proj(x)

@use_np
class CategoricalFeatureNet(HybridBlock):
    def __init__(self, num_class, out_units, cfg=None, prefix=None, params=None):
        super().__init__(prefix=prefix, params=params)
        if cfg is None:
            cfg = CategoricalFeatureNet.get_cfg()
        else:
            cfg = CategoricalFeatureNet.get_cfg().clone_merge(cfg)
        self.cfg = cfg
        embed_initializer = mx.init.create(*cfg.initializer.embed)
        weight_initializer = mx.init.create(*cfg.initializer.weight)
        bias_initializer = mx.init.create(*cfg.initializer.bias)
        with self.name_scope():
            self.embedding = nn.Embedding(input_dim=num_class,
                                          output_dim=cfg.emb_units,
                                          weight_initializer=embed_initializer)
            self.proj = BasicMLP(in_units=cfg.emb_units,
                                 mid_units=cfg.mid_units,
                                 out_units=out_units,
                                 num_layers=cfg.num_layers,
                                 normalization=cfg.normalization,
                                 norm_eps=cfg.norm_eps,
                                 data_dropout=cfg.data_dropout,
                                 dropout=cfg.dropout,
                                 activation=cfg.activation,
                                 weight_initializer=weight_initializer,
                                 bias_initializer=bias_initializer)

    @staticmethod
    def get_cfg(key=None):
        if key is None:
            cfg = CfgNode()
            cfg.emb_units = 32
            cfg.mid_units = 64
            cfg.num_layers = 1
            cfg.data_dropout = False
            cfg.dropout = 0.1
            cfg.activation = 'leaky'
            cfg.normalization = 'layer_norm'
            cfg.norm_eps = 1e-5
            cfg.initializer = CfgNode()
            cfg.initializer.embed = ['xavier', 'gaussian', 'in', 1.0]
            cfg.initializer.weight = ['xavier', 'uniform', 'avg', 3.0]
            cfg.initializer.bias = ['zeros']
            return cfg
        else:
            raise NotImplementedError

    def hybrid_forward(self, F, feature):
        embed = self.embedding(feature)
        return self.proj(embed)

@use_np
class NumericalFeatureNet(HybridBlock):
    def __init__(self, input_shape, out_units, cfg=None, prefix=None, params=None):
        super().__init__(prefix=prefix, params=params)
        if cfg is None:
            cfg = NumericalFeatureNet.get_cfg()
        self.input_shape = input_shape
        self.need_first_reshape = isinstance(input_shape, (list, tuple)) and len(input_shape) != 1
        self.in_units = int(np.prod(input_shape))
        self.cfg = NumericalFeatureNet.get_cfg().clone_merge(cfg)
        weight_initializer = mx.init.create(*cfg.initializer.weight)
        bias_initializer = mx.init.create(*cfg.initializer.bias)
        with self.name_scope():
            if self.cfg.input_centering:
                self.data_bn = nn.BatchNorm(in_channels=self.in_units)
            self.proj = BasicMLP(in_units=self.in_units,
                                 mid_units=cfg.mid_units,
                                 out_units=out_units,
                                 num_layers=cfg.num_layers,
                                 normalization=cfg.normalization,
                                 norm_eps=cfg.norm_eps,
                                 data_dropout=cfg.data_dropout,
                                 dropout=cfg.dropout,
                                 activation=cfg.activation,
                                 weight_initializer=weight_initializer,
                                 bias_initializer=bias_initializer)

    @staticmethod
    def get_cfg(key=None):
        if key is None:
            cfg = CfgNode()
            cfg.input_centering = False
            cfg.mid_units = 128
            cfg.num_layers = 1
            cfg.data_dropout = False
            cfg.dropout = 0.1
            cfg.activation = 'leaky'
            cfg.normalization = 'layer_norm'
            cfg.norm_eps = 1e-5
            cfg.initializer = CfgNode()
            cfg.initializer.weight = ['xavier', 'uniform', 'avg', 3.0]
            cfg.initializer.bias = ['zeros']
        else:
            raise NotImplementedError
        return cfg

    def hybrid_forward(self, F, features):
        if self.need_first_reshape:
            features = F.np.reshape(features, (-1, self.in_units))
        if self.cfg.input_centering:
            features = self.data_bn(features)
        return self.proj(features)

@use_np
class FeatureAggregator(HybridBlock):
    def __init__(self, num_fields, out_shape, in_units,
                 cfg=None, prefix=None, params=None):
        super().__init__(prefix=prefix, params=params)
        if cfg is None:
            cfg = FeatureAggregator.get_cfg()
        self.cfg = FeatureAggregator.get_cfg().clone_merge(cfg)
        self.num_fields = num_fields
        if isinstance(out_shape, list):
            out_shape = tuple(out_shape)
        self.out_shape = out_shape
        self.in_units = in_units
        weight_initializer = mx.init.create(*self.cfg.initializer.weight)
        bias_initializer = mx.init.create(*self.cfg.initializer.bias)
        out_units = int(np.prod(out_shape))
        with self.name_scope():
            if self.cfg.agg_type == 'mean':
                in_units = in_units
            elif self.cfg.agg_type == 'concat':
                in_units = in_units * num_fields
            else:
                raise NotImplementedError
            mid_units = in_units if cfg.mid_units < 0 else cfg.mid_units
            self.proj = BasicMLP(in_units=in_units,
                                 mid_units=mid_units,
                                 out_units=out_units,
                                 num_layers=cfg.num_layers,
                                 normalization=cfg.normalization,
                                 norm_eps=cfg.norm_eps,
                                 dropout=cfg.dropout,
                                 data_dropout=cfg.data_dropout,
                                 activation=cfg.activation,
                                 weight_initializer=weight_initializer,
                                 bias_initializer=bias_initializer)

    @staticmethod
    def get_cfg(key=None):
        if key is None:
            cfg = CfgNode()
            cfg.agg_type = 'concat'
            cfg.mid_units = -1
            cfg.num_layers = 0
            cfg.data_dropout = False
            cfg.dropout = 0.1
            cfg.activation = 'tanh'
            cfg.normalization = 'layer_norm'
            cfg.norm_eps = 1e-5
            cfg.initializer = CfgNode()
            cfg.initializer.weight = ['xavier', 'uniform', 'avg', 3.0]
            cfg.initializer.bias = ['zeros']
        else:
            raise NotImplementedError
        return cfg

    def hybrid_forward(self, F, field_proj_features):
        """

        Parameters
        ----------
        field_proj_features
            List of projection features. All elements must have the same shape.

        Returns
        -------
        scores
            Shape (batch_size,) + out_shape
        """
        if len(field_proj_features) == 0:
            agg_features = field_proj_features[0]
        else:
            if self.cfg.agg_type == 'mean':
                agg_features = F.np.stack(field_proj_features)
                agg_features = F.np.mean(agg_features, axis=0)
            elif self.cfg.agg_type == 'concat':
                agg_features = F.np.concatenate(field_proj_features, axis=-1)
            else:
                # TODO(sxjscience) May try to implement more advanced pooling methods for
                #  multimodal data.
                raise NotImplementedError
        scores = self.proj(agg_features)
        if len(self.out_shape) != 1:
            scores = F.np.reshape(scores, (-1,) + self.out_shape)
        return scores

@use_np
class BERTForTabularBasicV1(HybridBlock):
    """The basic model for tabular classification + regression with
    BERT (and its variants like ALBERT, MobileBERT, ELECTRA, etc.)
    as the backbone for handling text data.

    Here, we use the backbone network to extract the contextual embeddings and use
    another dense layer to map the contextual embeddings to the class scores.

    Input:

    TextField + EntityField --> TextNet -------> TextFeature
        ...
    CategoricalField --> CategoricalNet --> CategoricalFeature  ==> AggregateNet --> logits/scores
        ...
    NumericalField ----> NumericalNet ----> NumericalFeature
    """
    def __init__(self, text_backbone,
                 feature_field_info,
                 label_shape=None,
                 cfg=None,
                 prefix=None,
                 params=None):
        """

        Parameters
        ----------
        text_backbone
            Backbone network for handling the text data
        feature_field_info
            The field information of the training data. Each will be a tuple:
            - (field_type, attributes)
        label_shape
            The shape of the label/number of classes. If we need a scalar, it will be an empty tuple "()".
        cfg
            The configuration of the network
        prefix
        params
        """
        super().__init__(prefix=prefix, params=params)
        self.cfg = BERTForTabularBasicV1.get_cfg()
        if cfg is not None:
            self.cfg = self.cfg.clone_merge(cfg)
        assert self.cfg.text_net.pool_type == 'cls'
        feature_units = self.cfg.feature_units
        if feature_units == -1:
            feature_units = text_backbone.units
        if isinstance(label_shape, int):
            out_shape = (label_shape,)
        elif label_shape is None:
            out_shape = ()
        else:
            out_shape = label_shape
        with self.name_scope():
            self.text_backbone = text_backbone
            self.feature_field_info = feature_field_info
            self.categorical_fields = []
            self.numerical_fields = []
            self.categorical_networks = None
            self.numerical_network = None
            numerical_elements = None
            num_features = 1
            for i, (field_type_code, field_attrs) in enumerate(self.feature_field_info):
                if field_type_code == _C.CATEGORICAL:
                    if self.categorical_networks is None:
                        self.categorical_networks = nn.HybridSequential()
                    with self.categorical_networks.name_scope():
                        self.categorical_networks.add(
                            CategoricalFeatureNet(num_class=field_attrs['prop'].num_class,
                                                  out_units=feature_units,
                                                  cfg=cfg.categorical_net))
                        num_features += 1
                elif field_type_code == _C.NUMERICAL:
                    if numerical_elements is None:
                        numerical_elements = int(np.prod(field_attrs['prop'].shape))
                    else:
                        numerical_elements += int(np.prod(field_attrs['prop'].shape))
            if numerical_elements is not None:
                self.numerical_network = NumericalFeatureNet(input_shape=(numerical_elements,),
                                                             out_units=feature_units,
                                                             cfg=cfg.numerical_net)
                num_features += 1
            else:
                self.numerical_network = None
            self.agg_layer = FeatureAggregator(num_fields=num_features,
                                               out_shape=out_shape,
                                               in_units=feature_units,
                                               cfg=cfg.agg_net)

    @staticmethod
    def get_cfg(key=None):
        if key is None:
            cfg = CfgNode()
            cfg.feature_units = -1  # -1 means not given and we will use the units of BERT
            # TODO(sxjscience) Use a class to store the TextNet
            cfg.text_net = CfgNode()
            cfg.text_net.use_segment_id = True
            cfg.text_net.pool_type = 'cls'
            cfg.agg_net = FeatureAggregator.get_cfg()
            cfg.categorical_net = CategoricalFeatureNet.get_cfg()
            cfg.numerical_net = NumericalFeatureNet.get_cfg()
            cfg.initializer = CfgNode()
            cfg.initializer.weight = ['truncnorm', 0, 0.02]
            cfg.initializer.bias = ['zeros']
            return cfg
        else:
            raise NotImplementedError

    def initialize_with_pretrained_backbone(self, backbone_params_path, ctx=None):
        self.text_backbone.load_parameters(backbone_params_path, ctx=ctx)
        self.agg_layer.initialize(ctx=ctx)
        if self.categorical_networks is not None:
            self.categorical_networks.initialize(ctx=ctx)
        if self.numerical_network is not None:
            self.numerical_network.initialize(ctx=ctx)

    def hybrid_forward(self, F, features):
        """

        Parameters
        ----------
        features
            A list of field data

        Returns
        -------
        logits_or_scores
            Shape (batch_size,) + out_shape
        """
        field_features = []
        text_contextual_features = dict()
        categorical_count = 0
        numerical_samples = []
        for i, (field_type_code, field_attrs) in enumerate(self.feature_field_info):
            if field_type_code == _C.TEXT:
                batch_token_ids, batch_valid_length, batch_segment_ids, _ = features[i]
                if self.cfg.text_net.use_segment_id:
                    contextual_embedding, pooled_output = self.text_backbone(batch_token_ids,
                                                                             batch_segment_ids,
                                                                             batch_valid_length)
                else:
                    contextual_embedding = self.text_backbone(batch_token_ids, batch_valid_length)
                    pooled_output = contextual_embedding[:, 0, :]
                text_contextual_features[i] = contextual_embedding
                field_features.append(pooled_output)
            elif field_type_code == _C.ENTITY:
                # TODO Implement via segment-pool
                raise NotImplementedError('Currently not supported')
            elif field_type_code == _C.CATEGORICAL:
                batch_sample = features[i]
                extracted_feature = self.categorical_networks[categorical_count](batch_sample)
                categorical_count += 1
                field_features.append(extracted_feature)
            elif field_type_code == _C.NUMERICAL:
                batch_sample = features[i]
                numerical_samples.append(F.np.reshape(
                    batch_sample, (-1, int(np.prod(field_attrs['prop'].shape)))))
        if len(numerical_samples) > 0:
            if len(numerical_samples) == 1:
                numerical_feature = self.numerical_network(numerical_samples[0])
            else:
                numerical_feature = self.numerical_network(F.np.concatenate(numerical_samples,
                                                                            axis=-1))
            field_features.append(numerical_feature)
        return self.agg_layer(field_features)

max_stars_repo_name: RuohanW/autogluon

max_issues_repo_name: RuohanW/autogluon

max_forks_repo_name: RuohanW/autogluon