python-katlas 0.0.1__py3-none-any.whl

katlas/dl.py

@@ -0,0 +1,355 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_DL.ipynb.
+
+# %% auto 0
+__all__ = ['def_device', 'seed_everything', 'GeneralDataset', 'get_sampler', 'MLP_1', 'CNN1D_1', 'init_weights', 'lin_wn',
+           'conv_wn', 'CNN1D_2', 'train_dl', 'train_dl_cv', 'predict_dl']
+
+# %% ../nbs/04_DL.ipynb 4
+from fastbook import *
+import fastcore.all as fc,torch.nn.init as init
+from fastai.callback.training import GradientClip
+from torch.utils.data import WeightedRandomSampler
+
+# katlas
+from .core import Data
+from .feature import *
+from .train import *
+
+# sklearn
+from sklearn.model_selection import *
+from sklearn.metrics import mean_squared_error
+from scipy.stats import spearmanr,pearsonr
+
+# %% ../nbs/04_DL.ipynb 6
+def seed_everything(seed=123):
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+# %% ../nbs/04_DL.ipynb 8
+def_device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# %% ../nbs/04_DL.ipynb 13
+class GeneralDataset:
+    def __init__(self, 
+                 df, # a dataframe of values
+                 feat_col, # feature columns
+                 target_col=None # Will return test set for prediction if target col is None
+                ):
+        "A general dataset that can be applied to any dataframe"
+        
+        self.test = False if target_col is not None else True
+        
+        self.X = df[feat_col].values 
+        self.y = df[target_col].values if not self.test else None
+        
+        self.len = df.shape[0]
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, index):
+        X = torch.Tensor(self.X[index])
+        if self.test:
+            return X
+        else:
+            y = torch.Tensor(self.y[index])
+            return X, y
+
+# %% ../nbs/04_DL.ipynb 17
+def get_sampler(info,col):
+    
+    "For imbalanced data, get higher weights for less-represented samples"
+    
+    # get value counts
+    group_counts = info[col].value_counts()
+    
+    # to reduce the difference through log
+    # group_counts = group_counts.apply(lambda x: np.log(x+1.01))
+    
+    weights = 1. / group_counts[info[col]]
+
+    sample_weights = torch.from_numpy(weights.to_numpy())
+    sample_weights = torch.clamp_min(sample_weights,0.01)
+
+    sampler = WeightedRandomSampler(sample_weights, len(sample_weights),replacement=True)
+    
+    return sampler
+
+# %% ../nbs/04_DL.ipynb 23
+def MLP_1(num_features, 
+          num_targets,
+          hidden_units = [512, 218],
+          dp = 0.2):
+    
+    # Start with the first layer from num_features to the first hidden layer
+    layers = [
+        nn.Linear(num_features, hidden_units[0]),
+        nn.BatchNorm1d(hidden_units[0]),
+        nn.Dropout(dp),
+        nn.PReLU()
+    ]
+    
+    # Loop over hidden units to create intermediate layers
+    for i in range(len(hidden_units) - 1):
+        layers.extend([
+            nn.Linear(hidden_units[i], hidden_units[i+1]),
+            nn.BatchNorm1d(hidden_units[i+1]),
+            nn.Dropout(dp),
+            nn.PReLU()
+        ])
+    
+    # Add the output layer
+    layers.append(nn.Linear(hidden_units[-1], num_targets))
+    
+    model = nn.Sequential(*layers)
+    
+    return model
+
+# %% ../nbs/04_DL.ipynb 29
+class CNN1D_1(Module):
+    
+    def __init__(self, 
+                 num_features, # this does not matter, just for format
+                 num_targets):
+
+        self.conv1 = nn.Conv1d(in_channels=1, out_channels=3, kernel_size=3, dilation=1, padding=1, stride=1)
+        self.pool1 = nn.MaxPool1d(kernel_size=2, stride=2)
+        self.conv2 = nn.Conv1d(in_channels=3, out_channels=8, kernel_size=3, dilation=1, padding=1, stride=1)
+        self.pool2 = nn.MaxPool1d(kernel_size=2, stride=2)
+        self.flatten = Flatten()
+        self.fc1 = nn.Linear(in_features = int(8 * num_features/4), out_features=128)
+        self.fc2 = nn.Linear(in_features=128, out_features=num_targets)
+
+    def forward(self, x):
+        x = x.unsqueeze(1) # need shape (bs, 1, num_features) for CNN
+        x = self.pool1(nn.functional.relu(self.conv1(x)))
+        x = self.pool2(nn.functional.relu(self.conv2(x)))
+        # x = torch.flatten(x, 1)
+        x = self.flatten(x)
+        x = nn.functional.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+# %% ../nbs/04_DL.ipynb 33
+def init_weights(m, leaky=0.):
+    "Initiate any Conv layer with Kaiming norm."
+    if isinstance(m, (nn.Conv1d,nn.Conv2d,nn.Conv3d)): init.kaiming_normal_(m.weight, a=leaky)
+
+# %% ../nbs/04_DL.ipynb 34
+def lin_wn(ni,nf,dp=0.1,act=nn.SiLU):
+    "Weight norm of linear."
+    layers =  nn.Sequential(
+            nn.BatchNorm1d(ni),
+            nn.Dropout(dp),
+            nn.utils.weight_norm(nn.Linear(ni, nf)) )
+    if act: layers.append(act())
+    return layers
+
+# %% ../nbs/04_DL.ipynb 35
+def conv_wn(ni, nf, ks=3, stride=1, padding=1, dp=0.1,act=nn.ReLU):
+    "Weight norm of conv."
+    layers =  nn.Sequential(
+        nn.BatchNorm1d(ni),
+        nn.Dropout(dp),
+        nn.utils.weight_norm(nn.Conv1d(ni, nf, ks, stride, padding)) )
+    if act: layers.append(act())
+    return layers
+
+# %% ../nbs/04_DL.ipynb 36
+class CNN1D_2(nn.Module):
+    
+    def __init__(self, ni, nf, amp_scale = 16):
+        super().__init__()
+
+        cha_1,cha_2,cha_3 = 256,512,512
+        hidden_size = cha_1*amp_scale
+
+        cha_po_1 = hidden_size//(cha_1*2)
+        cha_po_2 = (hidden_size//(cha_1*4)) * cha_3
+        
+        self.lin = lin_wn(ni,hidden_size)
+        
+        # bs, 256, 16
+        self.view = View(-1,cha_1,amp_scale)
+        
+        self.conv1 = nn.Sequential(
+            conv_wn(cha_1, cha_2, ks=5, stride=1, padding=2, dp=0.1),
+            nn.AdaptiveAvgPool1d(output_size = cha_po_1),
+            conv_wn(cha_2, cha_2, ks=3, stride=1, padding=1, dp=0.1))
+        
+        self.conv2 = nn.Sequential(
+            conv_wn(cha_2, cha_2, ks=3, stride=1, padding=1, dp=0.3),
+            conv_wn(cha_2, cha_3, ks=5, stride=1, padding=2, dp=0.2))
+        
+        self.head = nn.Sequential(
+            nn.MaxPool1d(kernel_size=4, stride=2, padding=1),
+            nn.Flatten(),
+            lin_wn(cha_po_2,nf,act=None) )
+
+
+    def forward(self, x):
+        # amplify features to 4096
+        x = self.lin(x)
+        
+        # reshape to bs,256,16 for conv1d
+        x = self.view(x) 
+
+        x = self.conv1(x)
+        
+        x_s = x  # for skip connection (multiply)
+        x = self.conv2(x)
+        x = x * x_s
+
+        # Final block
+        x = self.head(x)
+
+        return x
+
+# %% ../nbs/04_DL.ipynb 40
+def train_dl(df, 
+            feat_col, 
+            target_col,
+            split, # tuple of numpy array for split index
+            model_func, # function to get pytorch model
+             n_epoch = 4, # number of epochs
+             bs = 32, # batch size
+             lr = 1e-2, # will be useless if lr_find is True
+            loss = mse, # loss function
+            save = None, # models/{save}.pth
+             sampler = None,
+             lr_find=False, # if true, will use lr from lr_find
+              ):
+    "A DL trainer."
+    
+    train = df.loc[split[0]]
+    valid = df.loc[split[1]]
+    
+    train_ds = GeneralDataset(train, feat_col, target_col)
+    valid_ds = GeneralDataset(valid, feat_col, target_col)
+    
+    n_workers = fc.defaults.cpus
+
+    if sampler is not None:
+        
+        train_dl = DataLoader(train_ds, batch_size=bs, sampler=sampler,num_workers=n_workers)
+        valid_dl = DataLoader(valid_ds, batch_size=bs, sampler=sampler,num_workers=n_workers)
+        
+        dls = DataLoaders(train_dl, valid_dl)
+        
+    else:
+        
+        dls = DataLoaders.from_dsets(train_ds, valid_ds, bs=bs, num_workers=n_workers)
+    
+    model = model_func()
+    
+    learn = Learner(dls.to(def_device), model.to(def_device), loss, 
+                    metrics= [PearsonCorrCoef(),SpearmanCorrCoef()],
+                    cbs = [GradientClip(1.0)] # prevent overfitting
+                   )
+    
+    if lr_find:
+        # get learning rate
+        lr = learn.lr_find()
+        plt.show()
+        plt.close()
+        print(lr)
+
+        
+    print('lr in training is', lr)
+    learn.fit_one_cycle(n_epoch,lr) #cbs = [SaveModelCallback(fname = 'best')] # save best model
+    
+    if save is not None:
+        learn.save(save)
+        
+    pred,target = learn.get_preds()
+    
+    pred = pd.DataFrame(pred.detach().cpu().numpy(),index=valid.index,columns=target_col)
+    target = pd.DataFrame(target.detach().cpu().numpy(),index=valid.index,columns=target_col)
+    
+    return target, pred
+
+# %% ../nbs/04_DL.ipynb 45
+@fc.delegates(train_dl)
+def train_dl_cv(df, 
+                feat_col, 
+                target_col, 
+                splits, # list of tuples
+                model_func, # functions like lambda x: return MLP_1(num_feat, num_target)
+                save:str=None,
+                **kwargs
+                ):
+    
+    OOF = []
+    metrics = []
+    
+    for fold,split in enumerate(splits):
+
+        print(f'------fold{fold}------')
+        
+        
+        fname=None
+        # save best model for each fold
+        if save is not None:
+            fname = f'{save}_fold{fold}'
+        
+        # train model
+        target, pred = train_dl(df,feat_col,target_col, split, model_func ,save=fname,**kwargs)
+
+        #------------get scores--------------
+        # get score metrics
+        mse, pearson_avg, _ = score_each(target,pred)
+        
+        # store metrics in a dictionary for the current fold
+        fold_metrics = {
+            'fold': fold,
+            'mse': mse,
+            'pearson_avg': pearson_avg
+        }
+        metrics.append(fold_metrics)
+
+        OOF.append(pred)
+        
+
+    # Concatenate OOF from each fold to a new dataframe
+    oof = pd.concat(OOF).sort_index()
+    
+    # Get metrics into a dataframe
+    metrics = pd.DataFrame(metrics)
+    
+    return oof, metrics
+
+# %% ../nbs/04_DL.ipynb 53
+def predict_dl(df, 
+               feat_col, 
+               target_col,
+               model, # model architecture
+               model_pth, # only name, not with .pth
+              ):
+    
+    "Predict dataframe given a deep learning model"
+    
+    test_dset = GeneralDataset(df,feat_col)
+    test_dl = DataLoader(test_dset,bs=512)
+    
+    
+    learn = Learner(None, model.to(def_device), loss_func=1)
+    learn.load(model_pth)
+    
+    learn.model.eval()
+    
+    preds = []
+    for data in test_dl:
+        inputs = data.to(def_device)
+        outputs = learn.model(inputs) #learn.model(x).sigmoid().detach().cpu().numpy()
+
+        preds.append(outputs.detach().cpu().numpy())
+
+    preds = np.concatenate(preds)
+    preds = pd.DataFrame(preds,index=df.index,columns=target_col)
+
+    return preds

katlas/feature.py

@@ -0,0 +1,290 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_feature.ipynb.
+
+# %% auto 0
+__all__ = ['get_rdkit', 'get_morgan', 'get_esm', 'get_t5', 'get_t5_bfd', 'reduce_feature', 'remove_hi_corr', 'preprocess']
+
+# %% ../nbs/01_feature.ipynb 4
+from fastbook import *
+import torch,re,joblib,gc,esm
+from tqdm.notebook import tqdm; tqdm.pandas()
+from .core import Data
+
+# Rdkit
+from rdkit import Chem
+from rdkit.ML.Descriptors import MoleculeDescriptors
+from rdkit.Chem import Draw,Descriptors,AllChem
+
+# Models
+from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
+from fairscale.nn.wrap import enable_wrap, wrap
+from transformers import T5Tokenizer, T5EncoderModel, T5Model
+
+# Dimension Reduction
+from sklearn import set_config
+from sklearn.decomposition import PCA
+from sklearn.manifold import TSNE
+from sklearn.preprocessing import StandardScaler
+from umap.umap_ import UMAP
+
+set_config(transform_output="pandas")
+
+# %% ../nbs/01_feature.ipynb 7
+def get_rdkit(df: pd.DataFrame, # a dataframe that contains smiles
+              col:str = "SMILES", # colname of smile
+              normalize: bool = True, # normalize features using StandardScaler()
+              ):
+    "Extract chemical features from smiles via rdkit.Chem.Descriptors; if normalize, apply StandardScaler"
+    
+    mols = [Chem.MolFromSmiles(smi) for smi in df[col]]
+    desc_names = [desc_name[0] for desc_name in Descriptors.descList]
+    desc_calc = MoleculeDescriptors.MolecularDescriptorCalculator(desc_names)
+    desc_values = [desc_calc.CalcDescriptors(mol) for mol in mols]
+    feature_df = pd.DataFrame(np.stack(desc_values), index=df.index,columns=desc_names)
+    
+    if normalize:
+        feature_df = StandardScaler().fit_transform(feature_df)
+        
+    # feature_df = feature_df.reset_index()
+    return feature_df
+
+# %% ../nbs/01_feature.ipynb 11
+def get_morgan(df: pd.DataFrame, # a dataframe that contains smiles
+               col: str = "SMILES", # colname of smile
+               radius=3
+              ):
+    "Get 2048 morgan fingerprint (binary feature) from smiles in a dataframe"
+    mols = [Chem.MolFromSmiles(smi) for smi in df[col]]
+    morgan_fps = [AllChem.GetMorganFingerprintAsBitVect(mol, radius=radius, nBits=2048) for mol in mols]
+    fp_df = pd.DataFrame(np.array(morgan_fps), index=df.index)
+    fp_df.columns = "morgan_" + fp_df.columns.astype(str)
+    return fp_df
+
+# %% ../nbs/01_feature.ipynb 15
+def get_esm(df:pd.DataFrame, # a dataframe that contains amino acid sequence
+            col: str = 'sequence', # colname of amino acid sequence
+            model_name: str = "esm2_t33_650M_UR50D", # Name of the ESM model to use for the embeddings.
+            ):
+    
+    "Extract esmfold2 embeddings from protein sequence in a dataframe"
+    
+    # Initialize distributed world with world_size 1
+    if not torch.distributed.is_initialized():
+        url = "tcp://localhost:23456"
+        torch.distributed.init_process_group(backend="nccl", init_method=url, world_size=1, rank=0)
+    
+    #get number of repr layers
+    match = re.search(r'_t(\d+)_', model_name)
+    number = int(match.group(1))
+    print(f"repr_layers number for model {model_name} is {number}.")
+    print("You can also choose other esm2 models:",
+          "\nesm2_t48_15B_UR50D\nesm2_t36_3B_UR50D\nesm2_t33_650M_UR50D\nesm2_t30_150M_UR50D\nesm2_t12_35M_UR50D\nesm2_t6_8M_UR50D\n")
+
+    # Download model data from the hub
+    model_data, regression_data = esm.pretrained._download_model_and_regression_data(model_name)
+
+    # Initialize the model with FSDP wrapper
+    fsdp_params = dict(
+        mixed_precision=True,
+        flatten_parameters=True,
+        state_dict_device=torch.device("cpu"),  # reduce GPU mem usage
+        cpu_offload=True,  # enable cpu offloading
+    )
+
+    with enable_wrap(wrapper_cls=FSDP, **fsdp_params):
+        model, vocab = esm.pretrained.load_model_and_alphabet_core(
+            model_name, model_data, regression_data
+        )
+        batch_converter = vocab.get_batch_converter()
+        model.eval()
+
+        # Wrap each layer in FSDP separately
+        for name, child in model.named_children():
+            if name == "layers":
+                for layer_name, layer in child.named_children():
+                    wrapped_layer = wrap(layer)
+                    setattr(child, layer_name, wrapped_layer)
+        model = wrap(model)
+
+        # Define the feature extraction function
+        def esm_embeddings(r, colname=col):
+            data = [('protein', r[colname])]
+            labels, strs, tokens = batch_converter(data)
+            with torch.no_grad():
+                results = model(tokens.cuda(), repr_layers=[number], return_contacts=False)
+            rpr = results["representations"][number].squeeze()
+            rpr = rpr[1 : len(r[colname]) + 1].mean(0).detach().cpu().numpy()
+
+            del results, labels, strs, tokens, data #especially need to delete those on cuda: tokens, results
+            gc.collect()
+
+            return rpr
+        
+        # Apply the feature extraction function to each row in the DataFrame
+        series = df.progress_apply(esm_embeddings, axis=1)
+        df_feature = pd.DataFrame(series.tolist(), index=df.index)
+        df_feature.columns = 'esm_' + df_feature.columns.astype(str)
+
+        return df_feature
+
+# %% ../nbs/01_feature.ipynb 19
+def get_t5(df: pd.DataFrame, 
+           col: str = 'sequence'
+           ):
+    "Extract ProtT5-XL-uniref50 embeddings from protein sequence in a dataframe"
+    
+    # Reference: https://github.com/agemagician/ProtTrans/tree/master/Embedding/PyTorch/Advanced
+    # Load the tokenizer
+    tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc', do_lower_case=False)
+
+    # Load the model
+    model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc").to('cuda')
+
+    # Set the model precision based on the device
+    model.half()
+    
+    def T5_embeddings(sequence):
+        seq_len = len(sequence)
+        # Prepare the protein sequences as a list
+        sequence = [" ".join(list(re.sub(r"[UZOB]", "X", sequence)))]
+
+        # Tokenize sequences and pad up to the longest sequence in the batch
+        ids = tokenizer.batch_encode_plus(sequence, add_special_tokens=True, padding="longest")
+        input_ids = torch.tensor(ids['input_ids']).to('cuda')
+        attention_mask = torch.tensor(ids['attention_mask']).to('cuda')
+
+        # Generate embeddings
+        with torch.no_grad():
+            embedding_rpr = model(input_ids=input_ids, attention_mask=attention_mask)
+
+        emb_mean = embedding_rpr.last_hidden_state[0][:seq_len].detach().cpu().numpy().mean(axis=0)
+
+        return emb_mean
+
+    series = df[col].progress_apply(T5_embeddings)
+        
+
+    T5_feature = pd.DataFrame(series.tolist(),index=df.index)
+    T5_feature.columns = 'T5_' + T5_feature.columns.astype(str)
+    
+    return T5_feature
+
+# %% ../nbs/01_feature.ipynb 22
+def get_t5_bfd(df:pd.DataFrame, 
+               col: str = 'sequence'
+               ):
+    
+    "Extract ProtT5-XL-BFD embeddings from protein sequence in a dataframe"
+    # Reference: https://github.com/agemagician/ProtTrans/tree/master/Embedding/PyTorch/Advanced
+    # Load the tokenizer
+    tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_bfd', do_lower_case=False)
+
+    model = T5Model.from_pretrained("Rostlab/prot_t5_xl_bfd").to('cuda')
+
+    model.eval()
+    
+    def T5_embeddings_bfd(sequence, device = 'cuda'):
+
+        seq_len = len(sequence)
+
+        # Prepare the protein sequences as a list
+        sequence = [" ".join(list(re.sub(r"[UZOB]", "X", sequence)))]
+
+        # Tokenize sequences and pad up to the longest sequence in the batch
+        ids = tokenizer.batch_encode_plus(sequence, add_special_tokens=True, padding="longest")
+        input_ids = torch.tensor(ids['input_ids']).to(device)
+        attention_mask = torch.tensor(ids['attention_mask']).to(device)
+
+        # Generate embeddings
+        with torch.no_grad():
+            embedding_rpr = model(input_ids=input_ids, attention_mask=attention_mask, decoder_input_ids = input_ids)
+
+        emb_mean = embedding_rpr.last_hidden_state[0][:seq_len].detach().cpu().numpy().mean(axis=0)
+
+        return emb_mean
+
+    series = df[col].progress_apply(T5_embeddings_bfd)
+        
+
+    T5_feature = pd.DataFrame(series.tolist(),index=df.index)
+    T5_feature.columns = 'T5bfd_' + T5_feature.columns.astype(str)
+    
+    return T5_feature
+
+# %% ../nbs/01_feature.ipynb 26
+def reduce_feature(df: pd.DataFrame, 
+                   method: str='pca', # dimensionality reduction method, accept both capital and lower case
+                   complexity: int=20, # None for PCA; perfplexity for TSNE, recommend: 30; n_neigbors for UMAP, recommend: 15
+                   n: int=2, # n_components
+                   load: str=None, # load a previous model, e.g. model.pkl
+                   save: str=None, # pkl file to be saved, e.g. pca_model.pkl
+                   seed: int=123, # seed for random_state
+                   **kwargs, # arguments from PCA, TSNE, or UMAP depends on which method to use
+                  ):
+    
+    "Reduce the dimensionality given a dataframe of values"
+    
+    method = method.lower()
+    assert method in ['pca','tsne','umap'], "Please choose a method among PCA, TSNE, and UMAP"
+    
+    if load is not None:
+        reducer = joblib.load(load)
+    else:
+        if method == 'pca':
+            reducer = PCA(n_components=n, random_state=seed,**kwargs)
+        elif method == 'tsne':
+            reducer = TSNE(n_components=n,
+                           random_state=seed, 
+                           perplexity = complexity, # default from official is 30 
+                          **kwargs)
+        elif method == 'umap':
+            reducer = UMAP(n_components=n, 
+                           random_state=seed, 
+                           n_neighbors=complexity, # default from official is 15, try 15-200
+                          **kwargs)
+        else:
+            raise ValueError('Invalid method specified')
+
+    proj = reducer.fit_transform(df)
+    embedding_df = pd.DataFrame(proj).set_index(df.index)
+    embedding_df.columns = [f"{method.upper()}{i}" for i in range(1, n + 1)]
+
+    if save is not None:
+        path = Path(save)
+        path.parent.mkdir(exist_ok=True)
+        
+        joblib.dump(reducer, save)
+
+    return embedding_df
+
+# %% ../nbs/01_feature.ipynb 29
+def remove_hi_corr(df: pd.DataFrame, 
+                   thr: float=0.98 # threshold
+                   ):
+    "Remove highly correlated features in a dataframe given a pearson threshold"
+    
+    # Create correlation matrix
+    corr_matrix = df.corr().abs()
+    
+    # Select upper triangle of correlation matrix
+    upper = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(bool))
+    
+    # Find index of feature columns with correlation greater than threshold
+    to_drop = [column for column in upper.columns if any(upper[column] > thr)]
+    
+    # Drop features 
+    df = df.drop(to_drop, axis=1)
+    
+    return df
+
+# %% ../nbs/01_feature.ipynb 33
+def preprocess(df: pd.DataFrame,
+               thr: float=0.98):
+    
+    "Remove features with no variance, and highly correlated features based on threshold"
+    
+    df_original = df.copy()
+    df = df.loc[:,df.std() != 0]
+    df = remove_hi_corr(df, thr)
+    dropping_col = set(df_original.columns) - set(df.columns)
+    print(f'removing columns: {dropping_col}')
+    return df

katlas/imports.py

@@ -0,0 +1,7 @@
+from katlas.core import *
+from katlas.feature import *
+from katlas.plot import *
+
+# for training
+from katlas.train import *
+from katlas.dl import *