A classification framework to enhance your habitat distribution models
import torch
import multiprocessing
import torchmetrics
import numpy as np
import scipy.sparse
import pickle
import sklearn.preprocessing
import captum
import tqdm
import os
import sys

script_directory = os.path.dirname(os.path.abspath('__file__'))
framework_directory = os.path.abspath(os.path.join(script_directory, '..'))
sys.path.append(framework_directory)

from utils import SparseDataset

from models.mlp import MLP
from data.preprocess_data import add_rank, add_fold_assignments, add_standardization

X = scipy.sparse.load_npz(os.path.join(framework_directory, 'Data/input_data.npz'))
y = np.load(os.path.join(framework_directory, 'Data/target_values.npy'))
with open(os.path.join(framework_directory, 'Data/le_species.pkl'), 'rb') as f:
    le_species = pickle.load(f)
with open(os.path.join(framework_directory, 'Data/le_header.pkl'), 'rb') as f:
    le_header = pickle.load(f)
split_assignments = np.load(os.path.join(framework_directory, 'Data/split_assignments.npy'))

attributions = torch.zeros((2, len(le_header.classes_), X.shape[1]))
ablations = torch.zeros((2, len(le_header.classes_), 2))
ranks = torch.zeros((len(le_species.classes_))).cuda()
predictions = []
models = {}

X = add_rank(X, len(le_species.classes_))

for fold in range(2):
    print('\n' + '*'*11)
    print('* Fold: {} *'.format(fold))
    print('*'*11 + '\n')
    
    X_train, X_test, y_train, y_test = add_fold_assignments(X, y, split_assignments, fold)
    
    X_test_species = X_test[:, :len(le_species.classes_)]

    sorted_indices = []

    for i in range(X_test_species.shape[0]):
        data = X_test_species[i].data
        col_indices = X_test_species[i].indices
        sorted_idx = col_indices[np.argsort(data)][::-1]
        sorted_indices.append(torch.from_numpy(sorted_idx.copy()).cuda().long())

    X_train, scaler = add_standardization(X_train, scaler=None)
    X_test, scaler = add_standardization(X_test, scaler=scaler)
    
    train_dataloader = torch.utils.data.DataLoader(dataset=SparseDataset(X_train, y_train), batch_size=256, shuffle=True, num_workers=multiprocessing.cpu_count())
    test_dataloader = torch.utils.data.DataLoader(dataset=SparseDataset(X_test, y_test), batch_size=256, shuffle=False, num_workers=multiprocessing.cpu_count())

    model = MLP(input_dim=X.shape[1], output_dim=len(le_header.classes_))
    model = torch.nn.DataParallel(model)
    model.cuda()
    
    criterion = torch.nn.CrossEntropyLoss()
    criterion.cuda()
    
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0001)
    
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.1, patience=5)
    
    best_fold_accuracy = -np.inf
    best_epoch = -1
    no_improvement_epochs = 0
    
    for epoch in range(1):
        
        for batch, labels in train_dataloader:
            batch = batch.requires_grad_().cuda()
            labels = labels.cuda()
            optimizer.zero_grad()
            outputs = model(batch)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()    
            
        metric = torchmetrics.Accuracy(task="multiclass", num_classes=len(le_header.classes_), average='micro', top_k=1)
        metric.cuda()
        
        model.eval()
        
        with torch.no_grad():
            for batch, labels in test_dataloader:
                batch = batch.cuda()
                labels = labels.cuda()
                outputs = model(batch)
                accuracy = metric(outputs, labels)
        accuracy_epoch = metric.compute() * 100
        print('-'*30)
        print('Epoch {} completed.'.format(epoch))
        print(f'Accuracy: {accuracy_epoch:.4f}%.')
        print('-'*30 + '\n')
        model.train()
        
        scheduler.step(accuracy)
        
        if accuracy_epoch > best_fold_accuracy:
            best_fold_accuracy = accuracy_epoch
            best_epoch = epoch
            no_improvement_epochs = 0
        else:
            no_improvement_epochs += 1
        if no_improvement_epochs == 10:             
            print(f'Early stopping occurred at epoch {epoch} with best epoch = {best_epoch} and best accuracy = {best_fold_accuracy:.4f}%\n')
            break
            
    if no_improvement_epochs != 10:
        print(f'Early stopping dit not occur as best epoch = {best_epoch} and best accuracy = {best_fold_accuracy:.4f}%\n')

    models[f'fold_{fold}'] = model.state_dict()

    with torch.no_grad():
        predictions_fold = model(torch.from_numpy(X_test.toarray()).float())

    predictions_fold = scipy.special.softmax(predictions_fold.cpu(), axis=1)
    predictions_fold = [np.argmax(sub_arr) for sub_arr in predictions_fold]
    
    predictions.append(np.asarray(predictions_fold))

    X_test_tensor = torch.from_numpy(X_test.toarray()).float().cuda()
    y_test_tensor = torch.from_numpy(y_test).cuda()

    num_samples = len(X_test_tensor)
    num_classes = len(le_header.classes_)
    batch_size = 256

    integrated_gradients = captum.attr.IntegratedGradients(model)
    feature_ablation = captum.attr.FeatureAblation(model) 

    feature_mask = torch.zeros((X_test.shape[1]), dtype=torch.int).cuda()
    feature_mask[:len(le_species.classes_)] = 0 
    feature_mask[len(le_species.classes_):] = 1

    attributions_sums = torch.zeros(num_classes, X_test_tensor.shape[1]).cuda()
    ablations_sums = torch.zeros(num_classes, 2).cuda()

    class_counts = torch.zeros(num_classes).cuda()
    
    for i in tqdm.tqdm(range(0, num_samples, batch_size), desc=f'Interpretability of fold {fold}'):
        start_idx = i
        end_idx = min(i + batch_size, num_samples)

        batch_X = X_test_tensor[start_idx:end_idx]
        batch_y = y_test_tensor[start_idx:end_idx]
        sorted_indices_batch = sorted_indices[start_idx:end_idx]
        
        attributions_batch = integrated_gradients.attribute(batch_X, target=batch_y)
        ablations_batch = feature_ablation.attribute(batch_X, target=batch_y, feature_mask=feature_mask)
       
        attributions_batch = torch.abs(attributions_batch)
        ablations_batch = torch.abs(ablations_batch)

        ablations_internal = ablations_batch[:, :len(le_species.classes_)]
        ablations_external = ablations_batch[:, len(le_species.classes_):]

        mean_internal = torch.mean(ablations_internal, dim=1)
        mean_external = torch.mean(ablations_external, dim=1)

        ablations_batch = torch.stack((mean_internal, mean_external), dim=1)

        for j in range(num_classes):
            mask = batch_y == j
            class_counts[j] += mask.sum()
            attributions_sums[j] += attributions_batch[mask].sum(dim=0)
            ablations_sums[j] += ablations_batch[mask].sum(dim=0)

        attributions_ranks = attributions_batch[:, :len(le_species.classes_)]
        
        for j in range(len(attributions_ranks)):
            result = torch.zeros(len(le_species.classes_)).cuda()
            result[:len(sorted_indices[j])] = attributions_ranks[j][sorted_indices[j]]
            ranks += result

    species_per_plot = np.diff(X_test[:, :len(le_species.classes_)].indptr)
    plots_per_species = torch.zeros((len(le_species.classes_),), dtype=torch.int).cuda()
    
    for species_count in species_per_plot:
        plots_per_species[:species_count + 1] += 1

    plots_per_species[plots_per_species == 0] = 1

    ranks = ranks / plots_per_species

    class_counts[class_counts == 0] = 1

    attributions_fold = attributions_sums / class_counts.view(-1, 1)
    ablations_fold = ablations_sums / class_counts.view(-1, 1)    

    attributions[fold] = attributions_fold
    ablations[fold] = ablations_fold

habitats = torch.zeros((num_classes))

for fold in range(2):
    for habitat in range(num_classes):
        if np.count_nonzero(y[split_assignments == fold] == habitat): 
            habitats[habitat] += 1

attributions_transposed = attributions.permute(0, 2, 1)
ablations_transposed = ablations.permute(0, 2, 1)

expanded_habitats = habitats.unsqueeze(0).unsqueeze(1)

averaged_attributions = attributions_transposed / expanded_habitats
averaged_ablations = ablations_transposed / expanded_habitats

attributions = averaged_attributions.mean(dim=0)
ablations = averaged_ablations.mean(dim=0)

attributions = attributions.permute(1, 0)
ablations = ablations.permute(1, 0)

torch.save(attributions, "attributions.pt")
torch.save(ablations, "ablations.pt")
torch.save(ranks, "ranks.pt")
with open('predictions.pkl', 'wb') as file:
    pickle.dump(predictions, file)
with open('models.pkl', 'wb') as file:
    pickle.dump(models, file)

Experiments/model_interpretability.py
cesar-leblanc/hdm-framework
from utils import bool_type

def add_all_parsers(parser):
    """
    Add all the necessary parsers based on the selected pipeline and model.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    _add_pipeline_parser(parser)  # Add the pipeline parser
    if parser.parse_known_args()[0].pipeline == 'check':  # Check if the selected pipeline is "check"
        _add_check_parser(parser)  # Add corresponding parser
    if parser.parse_known_args()[0].pipeline == 'dataset':  # Check if the selected pipeline is "dataset"
        _add_dataset_parser(parser)  # Add corresponding parser
    elif parser.parse_known_args()[0].pipeline == 'evaluation':  # Check if the selected pipeline is "evaluation"
        _add_evaluation_parser(parser)  # Add corresponding parser
    elif parser.parse_known_args()[0].pipeline == 'training':  # Check if the selected pipeline is "training"
        _add_training_parser(parser)  # Add corresponding parser
    elif parser.parse_known_args()[0].pipeline == 'prediction':  # Check if the selected pipeline is "prediction"
        _add_prediction_parser(parser)  # Add corresponding parser
    if parser.parse_known_args()[0].pipeline in ['evaluation', 'training', 'prediction']:  # Check if the selected pipeline is "evaluation", "tranining" or "prediction"
        if parser.parse_known_args()[0].model == 'mlp':  # Check if the selected model is MLP
            _add_mlp_parser(parser)  # Add corresponding parser
        elif parser.parse_known_args()[0].model == 'rfc':  # Check if the selected model is RFC
            _add_rfc_parser(parser)  # Add corresponding parser
        elif parser.parse_known_args()[0].model == 'xgb':  # Check if the selected model is XGB
            _add_xgb_parser(parser)  # Add corresponding parser
        elif parser.parse_known_args()[0].model == 'tnc':  # Check if the selected model is TNC
            _add_tnc_parser(parser)  # Add corresponding parser
        elif parser.parse_known_args()[0].model == 'ftt':  # Check if the selected model is FTT
            _add_ftt_parser(parser)  # Add corresponding parser
        _add_preprocessing_parser(parser)  # Add parser for preprocessing
    _add_filepaths_parser(parser)  # Add parser for filepaths
    _add_miscellaneous_parser(parser)  # Add miscellaneous parser

def _add_pipeline_parser(parser):
    """
    Add the pipeline parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_pipeline = parser.add_argument_group('Pipeline parameters')  # Create argument group for pipeline parameters
    group_pipeline.add_argument('--pipeline', required=True, choices=['check', 'dataset', 'evaluation', 'training', 'prediction'], help='The pipeline to choose.')

def _add_check_parser(parser):
    """
    Add the check parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_check = parser.add_argument_group('Check parameters')  # Create argument group for check parameters
    group_check.add_argument('--check_dependencies', type=bool_type, default=True, help='The choice to check if all the required dependencies are installed.')
    group_check.add_argument('--check_files', type=bool_type, default=True, help='The choice to check if all the necessary files and directories exist.')
    group_check.add_argument('--check_environment', type=bool_type, default=True, help='The choice to check if the environment is properly configured.')

def _add_dataset_parser(parser):
    """
    Add the dataset parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_dataset = parser.add_argument_group('Dataset parameters')  # Create argument group for dataset parameters
    group_dataset.add_argument('--level', default=3, choices=[1, 2, 3], help='The level of the EUNIS hierarchy.')
    group_dataset.add_argument('--min_year', type=int, default=0, help='The year from which we start keeping plots.')
    group_dataset.add_argument('--occurrences', type=int, default=10, help='The minimum number of occurrences for habitats and species.')
    group_dataset.add_argument('--countries', type=list, default=['all'], help='The list of countries to keep.')
    group_dataset.add_argument('--gbif_normalization', type=bool_type, default=True, help='The choice to normalize species names against the GBIF backbone.')
    group_dataset.add_argument('--folds', type=int, default=10, help='The number of folds for cross-validation.')
    group_dataset.add_argument('--east_west_bin_width_km', type=float, default=10.0, help='The width in kilometers for the east-west bins to use for the spatial split.')
    group_dataset.add_argument('--north_south_bin_width_km', type=float, default=10.0, help='The width in kilometers for the north-south bins to use for the spatial split.')
    group_dataset.add_argument('--location_features', type=bool_type, default=True, help='The choice to include several location-based features.')

def _add_evaluation_parser(parser):
    """
    Add the evaluation parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_evaluation = parser.add_argument_group('Evaluation parameters')  # Create argument group for evaluation parameters
    group_evaluation.add_argument('--model', default='mlp', choices=['mlp', 'rfc', 'ftt', 'xgb', 'tnc'], help='The classifier model to evaluate.')
    group_evaluation.add_argument('--metric', default='accuracy', choices=['accuracy'], help='The classification score.')
    group_evaluation.add_argument('--average', default='micro', choices=['micro', 'macro'], help='The reduction that is applied over labels.')
    group_evaluation.add_argument('--k', type=int, default=1, help='The number of most likely outcomes considered to find the correct label.')
    group_evaluation.add_argument('--n_folds', type=int, default=10, help='The number of folds.')

def _add_training_parser(parser):
    """
    Add the training parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_training = parser.add_argument_group('Training parameters')  # Create argument group for training parameters
    group_training.add_argument('--model', default='mlp', choices=['mlp', 'rfc', 'ftt', 'xgb', 'tnc'], help='The classifier model to train.')
    group_training.add_argument('--metric', default='accuracy', choices=['accuracy'], help='The classification score.')
    group_training.add_argument('--average', default='micro', choices=['micro', 'macro'], help='The reduction that is applied over labels.')
    group_training.add_argument('--k', type=int, default=1, help='The number of most likely outcomes considered to find the correct label.')

def _add_prediction_parser(parser):
    """
    Add the prediction parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_prediction = parser.add_argument_group('Prediction parameters')  # Create argument group for prediction parameters
    group_prediction.add_argument('--model', default='mlp', choices=['mlp', 'rfc', 'ftt', 'xgb', 'tnc'], help='The classifier model to predict.')
    group_prediction.add_argument('--gbif_normalization', type=bool_type, default=True, help='The choice to normalize species names against the GBIF backbone.')
    group_prediction.add_argument('--location_features', type=bool_type, default=True, help='The choice to include several location-based features.')

def _add_mlp_parser(parser):
    """
    Add the MLP parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_mlp = parser.add_argument_group('Multi-layer perceptron classifier parameters')  # Create argument group for MLP parameters
    group_mlp.add_argument('--num_hidden_layers', type=int, default=1, help='The number of hidden layers.')
    group_mlp.add_argument('--num_neurons', type=int, default=200, help='The hidden dimension (number of neurons / number of non-linear activation functions).')
    group_mlp.add_argument('--activation_func', default='relu', choices=['relu', 'sigmoid', 'tanh'], help='The activation function for the hidden layer.')
    group_mlp.add_argument('--optimizer', default='adam', choices=['sgd', 'adam'], help='The optimizer class.')
    group_mlp.add_argument('--learning_rate_init', type=float, default=0.001, help='The initial learning rate to use (controls the step-size in updating the weights).')
    group_mlp.add_argument('--batch_size', type=int, default=256, help='The size of minibatches.')
    group_mlp.add_argument('--criterion', default='cross-entropy', choices=['cross-entropy', 'bal-noised-top-k', 'imbal-noised-top-k'], help='The loss class.')
    group_mlp.add_argument('--weight_decay', type=float, default=0.0001, help='The strength of the L2 regularization term.')
    group_mlp.add_argument('--scheduler', default='plateau', choices=['plateau'], help='The step learning scheduler class.')
    group_mlp.add_argument('--num_epochs', type=int, default=100, help='The number of epochs (how many times each data point will be used).')
    group_mlp.add_argument('--num_iter_no_change', type=int, default=10, help='The patience parameter (maximum number of epochs to not improve).')
    group_mlp.add_argument('--momentum', type=float, default=0.9, help='The momentum for gradient descent update.')
    group_mlp.add_argument('--weight_init', default=None, choices=['kaiming', 'xavier'], help='The choice of weight initializations.') 
    group_mlp.add_argument('--dropout_p', type=float, default=0, help='The dropout probability.')
    group_mlp.add_argument('--use_batch_norm', type=bool_type, default=False, help='The choice to enable or disable batch normalization.')
    group_mlp.add_argument('--nesterov', type=bool_type, default=True, help="The choice to use Nesterov's momentum.")
    group_mlp.add_argument('--seed', type=int, default=123, help='The random seed for reproducibility.')
    group_mlp.add_argument('--factor', type=float, default=0.1, help='The factor by which the learning rate will be reduced.')
    group_mlp.add_argument('--patience', type=int, default=5, help='The number of epochs with no improvement after which learning rate will be reduced.')
    group_mlp.add_argument('--num_workers', type=int, default=-1, help='The number of subprocesses to use for data loading.')

def _add_rfc_parser(parser):
    """
    Add the RFC parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_rfc = parser.add_argument_group('Random forest classifier parameters')  # Create argument group for RFC parameters
    group_rfc.add_argument('--n_estimators', type=int, default=200, help='The number of trees in the forest.')
    group_rfc.add_argument('--criterion', default='gini', choices=['gini', 'entropy', 'log_loss'], help='The function to measure the quality of a split.')
    group_rfc.add_argument('--max_depth', type=int, default=None, help='The maximum depth of the tree.')
    group_rfc.add_argument('--min_samples_split', type=int, default=2, help='The minimum number of samples required to split an internal node.')
    group_rfc.add_argument('--min_samples_leaf', type=int, default=1, help='The minimum number of samples required to be at a leaf node.')
    group_rfc.add_argument('--min_weight_fraction_leaf', type=float, default=0.0, help='The minimum weighted fraction of the sum total of weights (of all the input samples) required to be at a leaf node.')
    group_rfc.add_argument('--max_features', default='sqrt', choices=['sqrt', 'log2', None], help='The number of features to consider when looking for the best split.')
    group_rfc.add_argument('--max_leaf_nodes', type=int, default=None, help='The value for which trees grow in best-first fashion.')
    group_rfc.add_argument('--min_impurity_decrease', type=float, default=0.0, help='The value for which a node will be split if this split induces a decrease of the impurity greater than or equal to this value.')
    group_rfc.add_argument('--bootstrap', type=bool_type, default=True, help='The choice to use bootstrap samples when building trees.')
    group_rfc.add_argument('--oob_score', type=bool_type, default=False, help='The choice to use out-of-bag samples to estimate the generalization score.')
    group_rfc.add_argument('--n_jobs', type=int, default=-1, help='The number of jobs to run in parallel.')
    group_rfc.add_argument('--random_state', type=int, default=123, help='The value that controls both the randomness of the bootstrapping of the samples used when building trees and the sampling of the features to consider when looking for the best split at each node.')
    group_rfc.add_argument('--verbose', default='n_estimators', choices=['n_estimators'] + list(map(str, range(0, 1000001))), help='The value that controls the verbosity when fitting and predicting.')
    group_rfc.add_argument('--warm_start', type=bool_type, default=False, help='The choice to reuse the solution of the previous call to fit and add mor estimators to the ensemble or to just fit a whole new forest.')
    group_rfc.add_argument('--class_weight', default=None, choices=['balanced', 'balanced_subsample', None], help='The weights associated with classes.')
    group_rfc.add_argument('--ccp_alpha', type=float, default=0.0, help='The complexity parameter used for Minimal Cost-Complexity Pruning.')
    group_rfc.add_argument('--max_samples', type=float, default=None, help='The number of samples to draw from X to train each base estimator if bootstrap is used.')

def _add_xgb_parser(parser):
    """
    Add the XGB parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_xgb = parser.add_argument_group('Extreme boosting classifier parameters')  # Create argument group for XGB parameters
    group_xgb.add_argument('--n_estimators', type=int, default=200, help='The number of boosting rounds.')
    group_xgb.add_argument('--max_depth', type=int, default=100, help='The maximum tree depth for base learners.')
    group_xgb.add_argument('--max_leaves', type=int, default=0, help='The maximum number of leaves; 0 indicates no limit.')
    group_xgb.add_argument('--max_bin', type=int, default=256, help='The maximum number of bins per feature if using histogram-based algorithm.')
    group_xgb.add_argument('--grow_policy', default='depthwise', choices=['depthwise', 'lossguide'], help='The tree growing policy.')
    group_xgb.add_argument('--learning_rate', type=float, default=0.1, help="The boosting learning rate (xgb's 'eta').")
    group_xgb.add_argument('--verbosity', type=int, default=1, help='The degree of verbosity.')
    group_xgb.add_argument('--objective', default='multi:softmax', choices=['multi:softmax', 'multi:softprob'], help='The learning task and the corresponding learning objective or a custom objective function to be used.')
    group_xgb.add_argument('--booster', default='gbtree', choices=['gbtree', 'gblinear', 'dart'], help='The booster to use.')
    group_xgb.add_argument('--tree_method', default='auto', choices=['auto', 'exact', 'approx', 'hist', 'gpu_hist'], help='The tree method to use.')
    group_xgb.add_argument('--n_jobs', type=int, default=-1, help='The number of parallel threads used to run xgboost.')
    group_xgb.add_argument('--gamma', type=float, default=0, help='The minimum loss reduction required to make a further partition on a leaf node of the tree.')
    group_xgb.add_argument('--min_child_weight', type=float, default=1, help='The minimum sum of instance weight (hessian) needed in a child.')
    group_xgb.add_argument('--max_delta_step', type=float, default=0, help="The maximum delta step we allow each tree's weight estimation to be.")
    group_xgb.add_argument('--subsample', type=float, default=0.8, help='The subsample ratio of the training instance.')
    group_xgb.add_argument('--sampling_method', default='uniform', choices=['uniform', 'gradient_based'], help='The sampling method.')
    group_xgb.add_argument('--colsample_bytree', type=float, default=1, help='The subsample ratio of columns when constructing each tree.')
    group_xgb.add_argument('--colsample_bylevel', type=float, default=1, help='The subsample ratio of columns for each level.')
    group_xgb.add_argument('--colsample_bynode', type=float, default=1, help='The subsample ratio of columns for each split.')
    group_xgb.add_argument('--reg_alpha', type=float, default=0, help="The L1 regularization term on weights (xgb's alpha).")
    group_xgb.add_argument('--reg_lambda', type=float, default=1, help="The L2 regularization term on weights (xgb's lambda).")
    group_xgb.add_argument('--scale_pos_weight', type=float, default=None, help='The balancing of positive and negative weights.')
    group_xgb.add_argument('--base_score', type=float, default=None, help='The initial prediction score of all instances, global bias.')
    group_xgb.add_argument('--random_state', type=int, default=123, help='The random number seed.')
    group_xgb.add_argument('--missing', type=float, default=float('nan'), help='The value in the data which needs to be present as a missing value.')
    group_xgb.add_argument('--num_parallel_tree', type=int, default=1, help='The value used for boosting random forest.')
    group_xgb.add_argument('--monotone_constraints', type=str, default=None, help='The constraint of variable monotonicity.')
    group_xgb.add_argument('--interaction_constraints', type=str, default=None, help='The constraints for interaction representing permitted interactions.')
    group_xgb.add_argument('--importance_type', type=str, default=None, help='The feature importance type for the feature importances property.')
    group_xgb.add_argument('--gpu_id', type=int, default=None, help='The device ordinal.')
    group_xgb.add_argument('--validate_parameters', type=bool_type, default=None, help='The choice to give warnings for unknown parameter.')
    group_xgb.add_argument('--predictor', default='auto', choices=['auto', 'cpu_predictor', 'gpu_predictor'], help='The value forces XGBoost to use specific predictor.')
    group_xgb.add_argument('--enable_categorical', type=bool_type, default=None, help='The support for categorical data.')
    group_xgb.add_argument('--feature_types', type=str, default=None, help='The value used for specifying feature types without constructing a dataframe.')
    group_xgb.add_argument('--max_cat_to_onehot', type=int, default=None, help='The threshold for deciding whether XGBoost should use one-hot encoding based split for categorical data.')
    group_xgb.add_argument('--max_cat_threshold', type=int, default=None, help='The maximum number of categories considered for each split.')
    group_xgb.add_argument('--eval_metric', type=str, default='merror', help='The metric used for monitoring the training result and early stopping.')
    group_xgb.add_argument('--early_stopping_rounds', type=int, default=10, help='The early stopping activation.')
    group_xgb.add_argument('--callbacks', type=list, default=None, help='The list of callback functions that are applied at end of each iteration.')
    group_xgb.add_argument('--use_label_encoder', type=bool_type, default=None, help='The choice to use label encoder.')

def _add_tnc_parser(parser):
    """
    Add the TNC parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_tnc = parser.add_argument_group('Tabnet classifier parameters')  # Create argument group for TNC parameters
    group_tnc.add_argument('--n_d', type=int, default=16, help='The width of the decision prediction layer.')
    group_tnc.add_argument('--n_a', type=int, default=16, help='The width of the attention embedding for each mask.')
    group_tnc.add_argument('--n_steps', type=int, default=3, help='The number of steps in the architecture (usually between 3 and 10).')
    group_tnc.add_argument('--gamma', type=float, default=1.3, help="The coefficient for feature reusage in the masks.")
    group_tnc.add_argument('--cat_idxs', type=list, default=[], help='The list of categorical features indices.')
    group_tnc.add_argument('--cat_dims', type=list, default=[], help='The list of categorical features number of modalities (number of unique values for a categorical feature).')
    group_tnc.add_argument('--cat_emb_dim', type=list, default=1, help='The list of embeddings size for each categorical features.')
    group_tnc.add_argument('--n_independent', type=int, default=1, help='The number of independent Gated Linear Units layers at each step.')
    group_tnc.add_argument('--n_shared', type=int, default=3, help='The number of shared Gated Linear Units at each step.')
    group_tnc.add_argument('--epsilon', type=float, default=1e-15, help='The value should be left untouched.')
    group_tnc.add_argument('--seed', type=int, default=123, help='The random seed for reproducibility.')
    group_tnc.add_argument('--momentum', type=float, default=0.02, help='The momentum for batch normalization.')
    group_tnc.add_argument('--clip_value', type=float, default=None, help='The gradient will be clipped at the given value.')
    group_tnc.add_argument('--lambda_sparse', type=float, default=1e-3, help='The extra sparsity loss coefficient.')
    group_tnc.add_argument('--optimizer_fn', default='adam', choices=['adam'], help='The PyTorch optimizer function.')
    group_tnc.add_argument('--optimizer_params', type=dict, default={'lr':0.02}, help='The parameters compatible with the PyTorch optimizer function used to initialize the optimizer.')
    group_tnc.add_argument('--scheduler_fn', default='step', choices=['step'], help='The PyTorch scheduler to change learning rates during training.')
    group_tnc.add_argument('--scheduler_params', type=dict, default={"step_size": 10, "gamma": 0.9}, help='The dictionnary of parameters to apply to the PyTorch scheduler.')
    group_tnc.add_argument('--model_name', type=str, default='TabNetClassifier', help='The name of the model used for saving in disk.')
    group_tnc.add_argument('--verbose', type=int, default=1, help='The verbosity for notebooks plots.')
    group_tnc.add_argument('--device_name', type=str, default='auto', help='The name of the device used for training.')
    group_tnc.add_argument('--mask_type', default='sparsemax', choices=['sparsemax', 'entmax'], help='The masking function to use for selection features.')
    group_tnc.add_argument('--grouped_features', type=list, default=[], help='The groups to allow the model to share its attention across features inside a same group.')
    group_tnc.add_argument('--n_shared_decoder', type=int, default=1, help='The number of shared GLU block in decoder.')
    group_tnc.add_argument('--n_indep_decoder', type=int, default=1, help='The number of independent GLU block in decoder.')
    group_tnc.add_argument('--eval_name', type=list, default=['eval'], help='The list of eval set names.')
    group_tnc.add_argument('--eval_metric', default=['accuracy'], choices=["accuracy", "balanced_accuracy", "logloss"], help='The list of evaluation metrics.')
    group_tnc.add_argument('--max_epochs', type=int, default=100, help='The maximum number of epochs for training.')
    group_tnc.add_argument('--patience', type=int, default=10, help='The number of consecutive epochs without improvement before performing early stopping.')
    group_tnc.add_argument('--weights', default=0, choices = [0, 1], help='The sampling parameter.')
    group_tnc.add_argument('--loss_fn', default='cross-entropy', choices=['cross-entropy'], help='The loss function for training.')    
    group_tnc.add_argument('--batch_size', type=int, default=1024, help='The number of examples per batch.')
    group_tnc.add_argument('--virtual_batch_size', type=int, default=128, help='The size of the mini batches used for "Ghost Batch Normalization".')
    group_tnc.add_argument('--num_workers', type=int, default=-1, help='The number or workers used in the data loader.')
    group_tnc.add_argument('--drop_last', type=bool_type, default=False, help='The choice to drop last batch if not complete during training.')
    group_tnc.add_argument('--callbacks', type=list, default=[], help='The list of custom callbacks.')
    group_tnc.add_argument('--pin_memory', type=bool_type, default=True, help='Whether to set pin_memory to True or False during training.')
    group_tnc.add_argument('--from_unsupervised', type=str, default=None, help='Use a previously self supervised model as starting weights.')
    group_tnc.add_argument('--augmentations', type=str, default=None, help='The choice to apply custom data augmentation pipeline during training.')
    group_tnc.add_argument('--warm_start', type=bool_type, default=False, help='The choice to fit twice the same model and start from a warm start.')
    group_tnc.add_argument('--compute_importance', type=bool_type, default=False, help='The choice to compute feature importance.')

def _add_ftt_parser(parser):
    """
    Add the FTT parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_ftt = parser.add_argument_group('Feature-tokenizer + transformer parameters')  # Create argument group for FTT parameters
    group_ftt.add_argument('--cat_cardinalities', type=list, default=None, help='The number of unique values for each feature.')
    group_ftt.add_argument('--d_token', type=int, default=16, help='The size of one token.')
    group_ftt.add_argument('--n_blocks', type=int, default=1, help='The number of Transformer blocks.')
    group_ftt.add_argument('--attention_n_heads', type=int, default=4, help='The number of attention heads.')
    group_ftt.add_argument('--attention_dropout', type=float, default=0.3, help='The dropout for attention blocks.')
    group_ftt.add_argument('--attention_initialization', type=str, default='kaiming', help='The initialization for attention blocks.')
    group_ftt.add_argument('--attention_normalization', type=str, default='LayerNorm', help='The normalization for attention blocks.')
    group_ftt.add_argument('--ffn_d_hidden', type=int, default=16, help='The input size for the second linear layer in the Feed-Forward Network module.')
    group_ftt.add_argument('--ffn_dropout', type=float, default=0.1, help='The dropout rate after the first linear layer in the Feed-Forward Network module.')
    group_ftt.add_argument('--ffn_activation', type=str, default='ReGLU', help='The activation used in the Feed-Forward Network.')
    group_ftt.add_argument('--ffn_normalization', type=str, default='LayerNorm', help='The normalization used in the Feed-Forward Network.')
    group_ftt.add_argument('--residual_dropout', type=float, default=0.0, help='The dropout rate rate for the output of each residual branch of all Transformer blocks.')
    group_ftt.add_argument('--prenormalization', type=bool_type, default=True, help='The choice to place normalizations at the beginning of each residual branch.')
    group_ftt.add_argument('--first_prenormalization', type=bool_type, default=False, help='The choice to keep the first normalization from the first Transformer layer.')
    group_ftt.add_argument('--last_layer_query_idx', type=list, default=[-1], help='The indices of tokens that should be processed by the last Transformer block.')
    group_ftt.add_argument('--kv_compression_ratio', type=float, default=0.004, help='The choice to apply a technique to speed up attention modules when the number of features is large.')
    group_ftt.add_argument('--kv_compression_sharing', default='headwise', choices=[None, 'headwise', 'key-value', 'layerwise'], help='Weight sharing policy for the technique to speed up attention modules when the number of features is large.')
    group_ftt.add_argument('--head_activation', type=str, default='ReLU', help='The activation used in the heads.')
    group_ftt.add_argument('--head_normalization', type=str, default='LayerNorm', help='The normalization used in the heads.')
    group_ftt.add_argument('--learning_rate_init', type=float, default=0.001, help='The initial learning rate to use (controls the step-size in updating the weights).')
    group_ftt.add_argument('--optimizer', default='adamw', choices=['adamw'], help='The optimizer class.')
    group_ftt.add_argument('--batch_size', type=int, default=512, help='The size of minibatches.')
    group_ftt.add_argument('--criterion', default='cross-entropy', choices=['cross-entropy'], help='The loss class.')
    group_ftt.add_argument('--weight_decay', type=float, default=0.0001, help='The strength of the L2 regularization term.')
    group_ftt.add_argument('--scheduler', default='plateau', choices=['plateau'], help='The step learning scheduler class.')
    group_ftt.add_argument('--num_epochs', type=int, default=100, help='The number of epochs (how many times each data point will be used).')
    group_ftt.add_argument('--num_iter_no_change', type=int, default=10, help='The patience parameter (maximum number of epochs to not improve).')    
    group_ftt.add_argument('--seed', type=int, default=123, help='The random seed for reproducibility.')
    group_ftt.add_argument('--factor', type=float, default=0.1, help='The factor by which the learning rate will be reduced.')
    group_ftt.add_argument('--patience', type=int, default=5, help='The number of epochs with no improvement after which learning rate will be reduced.')
    group_ftt.add_argument('--num_workers', type=int, default=-1, help='The number of subprocesses to use for data loading.')

def _add_preprocessing_parser(parser):
    """
    Add the preprocessing parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_preprocessing = parser.add_argument_group('Preprocessing parameters')  # Create argument group for preprocessing parameters
    group_preprocessing.add_argument('--standardization', type=bool_type, default=True, help='The choice to standardize features by removing the mean and scaling to unit variance.')
    group_preprocessing.add_argument('--binarization', type=bool_type, default=False, help='The choice to binarize data (set feature values to 0 or 1) according to a threshold.')
    group_preprocessing.add_argument('--dropout', type=float, default=0.0, help='The dropout probability for dropping out the features.')
    group_preprocessing.add_argument('--noise', type=bool_type, default=False, help='The choice to add noise with random samples drawn from a normal distribution.')
    group_preprocessing.add_argument('--normalization', type=bool_type, default=False, help='The choice to normalize samples individually to unit norm.')
    group_preprocessing.add_argument('--log', type=bool_type, default=False, help='The choice to apply the natural logarithm to the features.')
    group_preprocessing.add_argument('--rank', type=bool_type, default=False, help='The choice to rank the features.')
    group_preprocessing.add_argument('--augmentation', type=bool_type, default=False, help='The choice to apply custom data augmentation pipeline during training.')    
    group_preprocessing.add_argument('--rare', type=bool_type, default=False, help='The choice to keep only the rarest labels.')
    group_preprocessing.add_argument('--endangered', type=bool_type, default=False, help='The choice to keep only the endangered habitats.')    
    group_preprocessing.add_argument('--features', type=list, default=['all'], help='The list of features the model will see during fit.')

def _add_filepaths_parser(parser):
    """
    Add the filepaths parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_data = parser.add_argument_group('Filepaths parameters')  # Create argument group for filepaths parameters
    group_data.add_argument('--data_filepath', type=str, default='Data/', help='The filepath to load/save the data from/to a disk file.')
    group_data.add_argument('--datasets_filepath', type=str, default='Datasets/', help='The filepath to load/save the datasets from/to a disk file.')
    group_data.add_argument('--models_filepath', type=str, default='Models/', help='The filepath to load/save the model from/to a disk file.')

def _add_miscellaneous_parser(parser):
    """
    Add the miscellaneous parser to the argument parser.

    Args:
        parser (argparse.ArgumentParser): The argument parser object.
    """
    group_miscellaneous = parser.add_argument_group('Miscellaneous parameters')  # Create argument group for miscellaneous parameters
    group_miscellaneous.add_argument('--print_parameters', type=bool_type, default=True, help='The choice to print the list of parameters.')
    group_miscellaneous.add_argument('--global_seed', type=int, default=123, help='The seed for reproducible output from run to run.')
    group_miscellaneous.add_argument('--use_gpu', type=bool_type, default=True, help='The choice to use GPU for the model.')
    group_miscellaneous.add_argument('--data_parallelism', type=bool_type, default=True, help='The choice to implement data parallelism at the module level.')
    group_miscellaneous.add_argument('--write_bytecode', type=bool_type, default=False, help='The choice to able or disable bytecode generation.')
    group_miscellaneous.add_argument('--disable_warnings', type=bool_type, default=True, help='The choice to able or disable warnings.')

cli.py
cesar-leblanc/hdm-framework
import scipy.sparse
import numpy as np
import pickle
import pandas as pd
import geopandas as gpd
import shapely
import itertools
import rioxarray
import pyproj
import torch
import joblib

def get_input_data(args):
    """
    Load the input data from a compressed sparse matrix file.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        scipy.sparse._csr.csr_matrix: The input data.
    """
    input_data_filepath = args.data_filepath + 'input_data.npz'  # Filepath of the input data
    X = scipy.sparse.load_npz(input_data_filepath)  # Load the input data from the file
    return X

def get_target_values(args):
    """
    Load the target values from a NumPy file.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        numpy.ndarray: The target values.
    """
    target_values_filepath = args.data_filepath + 'target_values.npy'  # Filepath of the target values
    y = np.load(target_values_filepath)  # Load the target values from the file
    return y

def get_le_species(args):
    """
    Load the species label encoder from a pickle file.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        sklearn.preprocessing._label.LabelEncoder: The loaded species label encoder object.
    """
    le_species_filepath = args.data_filepath + 'le_species.pkl'  # Filepath of the species label encoder
    with open(le_species_filepath, 'rb') as f:  # Open the file in binary mode
        le_species = pickle.load(f)  # Load the species label encoder from the file
    return le_species

def get_le_header(args):
    """
    Load the header label encoder from a pickle file.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        sklearn.preprocessing._label.LabelEncoder: The loaded header label encoder object.
    """
    le_header_filepath = args.data_filepath + 'le_header.pkl'  # Filepath of the header label encoder
    with open(le_header_filepath, 'rb') as f:  # Open the file in binary mode
        le_header = pickle.load(f)  # Load the header label encoder from the file
    return le_header

def get_split_assignments(args):
    """
    Load the split assignments from a NumPy array file.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        numpy.ndarray: The loaded split assignments.
    """
    split_assignments_filepath = args.data_filepath + 'split_assignments.npy'  # Filepath of the split assignments
    split_assignments = np.load(split_assignments_filepath)  # Load the split assignments from the file
    return split_assignments

def get_eunis_habitats(args):
    """
    Load the EUNIS habitats from an Excel file and filter them based on specific criteria.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        list: The filtered EUNIS habitats.
    """
    eunis_habitats_filepath = args.datasets_filepath + 'eunis_habitats.xlsx'  # Filepath of the EUNIS habitats Excel file
    eunis_habitats = pd.read_excel(eunis_habitats_filepath)  # Read the Excel file into a pandas DataFrame
    eunis_habitats = eunis_habitats['EUNIS 2020 code'].values.tolist()  # Extract the 'EUNIS 2020 code' column as a list
    eunis_habitats = [habitat for habitat in eunis_habitats if (not habitat.startswith('MA2') and len(habitat) == args.level) or (habitat.startswith('MA2') and len(habitat) == args.level + 2)]  # Filter the habitats based on their code length and prefixes
    return eunis_habitats

def get_eva_header(args, eunis_habitats):
    """
    Load and preprocess the eva_header data, filtering and transforming it based on specific criteria.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.
        eunis_habitats (list): A list of EUNIS habitats to filter the eva_header data.

    Returns:
        pandas.core.frame.DataFrame: The preprocessed eva_header data.
    """
    eva_header_filepath = args.datasets_filepath + 'eva_header.csv'  # Construct the filepath to the eva_header.csv file
    header_columns = ['PlotObservationID', 'Cover abundance scale', 'Date of recording', 'Expert System', 'Longitude', 'Latitude']  # Define the column names of the eva_header dataframe
    header_types = {'PlotObservationID': int, 'Cover abundance scale': str, 'Date of recording': int, 'Expert System': str, 'Longitude': float, 'Latitude': float}  # Define the data types of the eva_header dataframe
    eva_header = pd.read_csv(eva_header_filepath, sep='\t', usecols=header_columns, dtype=str, na_values='?')  # Read in the eva_header.csv file as a pandas dataframe
    eva_header['Cover abundance scale'] = eva_header['Cover abundance scale'].fillna('0')  # Replace missing values in the 'Cover abundance scale' column with '0'
    eva_header['Date of recording'] = eva_header['Date of recording'].fillna('0')  # Replace missing values in the 'Date of recording' column with '0'
    eva_header = eva_header.dropna()  # Remove any rows containing missing values
    eva_header = eva_header[eva_header['Cover abundance scale'] != 'Presence/Absence']  # Remove any rows where the 'Cover abundance scale' column has a value of 'Presence/Absence'
    eva_header["Date of recording"] = eva_header["Date of recording"].str.replace(':', '', regex=True).apply(lambda x: int(x[-4:]))  # Extract the year from the 'Date of recording' column and convert it to an integer
    eva_header = eva_header.loc[eva_header['Date of recording'] >= args.min_year]  # Remove any rows where the year in the 'Date of recording' column is less than the minimum year specified by the 'args' argument
    split_eva_header = eva_header['Expert System'].str.split(',', expand=True)  # Create a new dataframe for each split value in the 'classes' column
    split_eva_header = split_eva_header.stack().reset_index(level=1, drop=True).rename('Expert System')  # Stack the resulting dataframes vertically and rename the columns
    eva_header = eva_header.drop('Expert System', axis=1).join(split_eva_header)  # Join the original dataframe with the split dataframe on the index
    eva_header = eva_header.loc[eva_header['Expert System'].isin(eunis_habitats)]  # Filter out rows where the 'Expert System' column is not in the list of EUNIS habitats
    counts = eva_header['Expert System'].value_counts()  # Count the number of occurrences of each unique value in the 'classes' column
    keep = counts[counts >= args.occurrences].index.tolist()  # Keep only the values that appear 10 times or more
    eva_header = eva_header[eva_header['Expert System'].isin(keep)]  # Filter the dataframe to keep only the rows with values in the 'Expert System' column that appear 10 times or more
    eva_header = eva_header.astype(header_types)  # Convert the data types of the dataframe columns to the specified types
    eva_header = eva_header.round({'Longitude': 4, 'Latitude': 4})  # Round the 'Longitude' and 'Latitude' columns
    eva_header = eva_header.drop(['Cover abundance scale', 'Date of recording'], axis=1)  # Drop 'Cover abundance scale' and 'Date of recording' columns
    return eva_header  # Return the filtered dataframe

def get_test_header(args):
    """
    Load and preprocess the test header data, filtering and transforming it based on specific criteria.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        pandas.core.frame.DataFrame: The preprocessed test header data.
    """
    test_header_filepath = args.datasets_filepath + 'test_header.csv'  # Filepath of the test header CSV file
    header_columns = ['PlotObservationID', 'Longitude', 'Latitude']  # Specify the columns to be read from the CSV file
    header_types = {'PlotObservationID': int, 'Longitude': float, 'Latitude': float}  # Specify the data types for each column
    test_header = pd.read_csv(test_header_filepath, sep='\t', usecols=header_columns, dtype=str, na_values='?')  # Read the CSV file into a pandas DataFrame, using specified column names, data types, and handling missing values
    test_header = test_header.dropna()  # Drop rows with missing values from the DataFrame
    test_header = test_header.astype(header_types)  # Convert the data types of columns in the DataFrame to the specified types
    test_header = test_header.round({'Longitude': 4, 'Latitude': 4})  # Round the values in the 'Longitude' and 'Latitude' columns to 4 decimal places
    return test_header

def get_eva_species(args, eva_header):
    """
    Load and preprocess the eva_species data, filtering and transforming it based on specific criteria.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.
        eva_header (pandas.core.frame.DataFrame): The preprocessed eva_header data.

    Returns:
        pandas.core.frame.DataFrame: The preprocessed eva_species data.
    """
    eva_species_filepath = args.datasets_filepath + 'eva_species.csv'  # Filepath of the eva_species CSV file
    species_columns = ['PlotObservationID', 'Matched concept', 'Cover %']  # Specify the columns to be read from the CSV file
    species_types = {'PlotObservationID': int, 'Matched concept': str, 'Cover %': float}  # Specify the data types for each column
    eva_species = pd.read_csv(eva_species_filepath, delimiter='\t', usecols=species_columns, dtype=str)  # Read the CSV file into a pandas DataFrame, using specified column names, data types, and delimiter
    eva_species = eva_species.dropna()  # Drop rows with missing values from the DataFrame
    eva_species = eva_species.astype(species_types)  # Convert the data types of columns in the DataFrame to the specified types
    eva_species = eva_species[eva_species['PlotObservationID'].isin(eva_header['PlotObservationID'])]  # Filter the eva_species DataFrame to keep only the rows with PlotObservationID present in the eva_header DataFrame
    eva_species = eva_species.groupby(['PlotObservationID', 'Matched concept']).sum().reset_index()  # Group the DataFrame by 'PlotObservationID' and 'Matched concept', and sum the 'Cover %' values
    eva_species = eva_species[eva_species["Cover %"] > 0]  # Filter the DataFrame to keep only the rows with 'Cover %' greater than 0
    counts = eva_species['Matched concept'].value_counts()  # Count the number of occurrences of each unique value in the 'Matched concept' column
    keep = counts[counts >= args.occurrences].index.tolist()  # Keep only the values that appear 'args.occurrences' times or more
    eva_species = eva_species[eva_species['Matched concept'].isin(keep)]  # Filter the dataframe to keep only the rows with values in the 'Matched concept' column that appear 10 times or more
    return eva_species

def get_test_species(args, test_header):
    """
    Load and preprocess the test_species data, filtering and transforming it based on specific criteria.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.
        test_header (pandas.core.frame.DataFrame): The preprocessed test_header data.

    Returns:
        pandas.core.frame.DataFrame: The preprocessed test_species.
    """
    test_species_filepath = args.datasets_filepath + 'test_species.csv'  # Filepath of the test_species CSV file
    species_columns = ['PlotObservationID', 'Matched concept', 'Cover %']  # Specify the columns to be read from the CSV file
    species_types = {'PlotObservationID': int, 'Matched concept': str, 'Cover %': float}  # Specify the data types for each column
    test_species = pd.read_csv(test_species_filepath, delimiter='\t', usecols=species_columns, dtype=str)  # Read the CSV file into a pandas DataFrame, using specified column names, data types, and delimiter
    test_species = test_species.dropna()  # Drop rows with missing values from the DataFrame
    test_species = test_species.astype(species_types)  # Convert the data types of columns in the DataFrame to the specified types
    test_species = test_species[test_species['PlotObservationID'].isin(test_header['PlotObservationID'])]  # Filter the test_species DataFrame to keep only the rows with PlotObservationID present in the test_header DataFrame
    test_species = test_species.groupby(['PlotObservationID', 'Matched concept']).sum().reset_index()  # Group the DataFrame by 'PlotObservationID' and 'Matched concept', and sum the 'Cover %' values
    test_species = test_species[test_species["Cover %"] > 0]  # Filter the DataFrame to keep only the rows with 'Cover %' greater than 0
    return test_species

def get_eva_to_gbif_species(args):
    """
    Load the eva_to_gbif_species data from a pickle file.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        dict: The loaded eva_to_gbif_species data.
    """
    eva_to_gbif_species_filepath = args.data_filepath + 'eva_to_gbif_species.pkl'  # Filepath of the eva_to_gbif_species pickle file
    with open(eva_to_gbif_species_filepath, 'rb') as f:  # Open the pickle file in binary mode for reading
        eva_to_gbif_species = pickle.load(f)  # Load the contents of the pickle file
    return eva_to_gbif_species

def get_altitudes(args):
    """
    Load and preprocess the digital elevation model (DEM) data.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        xarray.core.dataarray.DataArray: The processed DEM data.
    """
    altitudes_filepath = args.datasets_filepath + 'digital_elevation_model.tif'  # Construct the file path of the DEM file
    altitudes = rioxarray.open_rasterio(altitudes_filepath)  # Load the DEM data as an xarray.DataArray object
    altitudes = altitudes.rio.reproject("EPSG:3857")  # Reproject the DEM data to EPSG:3857 coordinate system
    altitudes.attrs['_FillValue'] = 0  # Set the fill value of the DEM to 0
    altitudes = altitudes.where(altitudes != 32767, 0)  # Replace nodata values with 0
    return altitudes  # Return the processed DEM data

def get_countries(args):
    """
    Load and preprocess the country geometries.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        geopandas.geodataframe.GeoDataFrame: The processed country geometries.
    """
    countries_columns = ['COUNTRY', 'geometry']  # Specify the columns to keep in the country geometries
    countries_filepath = args.datasets_filepath + 'world_countries.shp'  # Construct the file path of the country shapefile
    countries = gpd.read_file(countries_filepath)  # Load the country geometries as a GeoDataFrame object
    countries = countries[countries_columns]  # Select only the specified columns
    countries = countries.to_crs(crs=3857)  # Reproject the country geometries to EPSG:3857 coordinate system
    return countries  # Return the processed country geometries

def get_ecoregions(args):
    """
    Load and preprocess the ecoregion geometries.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        geopandas.geodataframe.GeoDataFrame: The processed ecoregion geometries.
    """
    ecoregions_columns = ['ECO_ID', 'geometry']  # Specify the columns to keep in the ecoregion geometries
    ecoregions_filepath = args.datasets_filepath + 'ecoregions.shp'  # Construct the file path of the ecoregion shapefile
    ecoregions = gpd.read_file(ecoregions_filepath)  # Load the ecoregion geometries as a GeoDataFrame object
    ecoregions = ecoregions[ecoregions_columns]  # Select only the specified columns
    ecoregions = ecoregions.to_crs(crs=3857)  # Reproject the ecoregion geometries to EPSG:3857 coordinate system
    return ecoregions  # Return the processed ecoregion geometries

def get_dunes(args):
    """
    Load and preprocess the dune geometries.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        shapely.geometry.multipolygon.MultiPolygon: The merged dune polygon.
    """
    dunes_columns = ['CODE', 'geometry']  # Specify the columns to keep in the dune geometries
    dunes_filepath = args.datasets_filepath + 'vegetation.shp'  # Construct the file path of the dune shapefile
    dunes = gpd.read_file(dunes_filepath)  # Load the dune geometries as a GeoDataFrame object
    dunes = dunes[dunes_columns]  # Select only the specified columns
    dunes = dunes[dunes['CODE'].isin([f'P{i}' for i in range(1,17)])]  # Filter the dunes to keep only those with CODE values matching P1 to P16
    dunes = dunes.to_crs(crs=3857)  # Reproject the dune geometries to EPSG:3857 coordinate system
    dunes = shapely.ops.unary_union(dunes['geometry'].tolist())  # Merge the dune geometries into a single polygon using unary union
    return dunes  # Return the merged dune polygon

def get_coasts(args):
    """
    Load and preprocess the coastal geometries.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        geopandas.geodataframe.GeoDataFrame: The dissolved coastal geometries.
    """
    categories = dict()  # Create an empty dictionary
    categories['ARC_COAST'] = ["White Sea", "Arctic Ocean", "Barentsz Sea", "Kara Sea", "Laptev Sea", "Norwegian Sea"]  # Select all the bodies of water that are part of the Arctic Ocean
    categories['ATL_COAST'] = ["Strait of Gibraltar", "Bristol Channel", "Irish Sea and St. George's Channel", "Inner Seas off the West Coast of Scotland", "North Atlantic Ocean", "Bay of Biscay", "Celtic Sea", "English Channel", "North Sea", "Skagerrak"]  # Select all the bodies of water that are part of the Atlantic Ocean
    categories['BAL_COAST'] = ["Gulf of Riga", "Baltic Sea", "Gulf of Finland", "Gulf of Bothnia", "Kattegat"]  # Select all the bodies of water that are part of the Baltic Sea
    categories['BLA_COAST'] = ["Black Sea", "Sea of Azov"]  # Select all the bodies of water that are part of the Black Sea
    categories['MED_COAST'] = ["Alboran Sea", "Ionian Sea", "Tyrrhenian Sea", "Adriatic Sea", "Mediterranean Sea - Eastern Basin", "Aegean Sea", "Sea of Marmara", "Balearic (Iberian Sea)", "Mediterranean Sea - Western Basin", "Ligurian Sea"]  # Select all the bodies of water that are part of the Mediterranean Sea
    coasts_columns = ['NAME', 'geometry']  # Specify the columns to keep in the coast geometries
    coasts_filepath = args.datasets_filepath + 'world_seas.shp'  # Construct the file path of the coast shapefile
    coasts = gpd.read_file(coasts_filepath)  # Load the coastal geometries as a GeoDataFrame object
    coasts = coasts[coasts_columns]  # Select only the specified columns
    coasts = coasts.to_crs(crs=3857)  # Reproject the coastal geometries to EPSG:3857 coordinate system
    coasts["geometry"] = coasts.boundary  # Extract the boundary of the coastal geometries
    for i in range(len(coasts)):  # Iterate through the bodies of water
        if coasts.loc[i, 'NAME'] in list(itertools.chain(*categories.values())):  # Check if the body of water is part of the Arctic Ocean, the Atlantic Ocean, the Baltic Sea, the Black Sea or the Mediterranean Sea
            coasts.loc[i, 'NAME'] = list({key: value for key, value in categories.items() if coasts.loc[i, 'NAME'] in value}.keys())[0]  # Replace the body of water by its larger part
    coasts = coasts.loc[coasts.loc[:, 'NAME'].isin(list(categories.keys()))]  # Filter the coastal geometries to keep only those with valid names
    coasts = coasts.dissolve(by='NAME')  # Dissolve the coastal geometries based on the 'NAME' column
    return coasts

def get_eunis_red_list_crosswalks(args):
    """
    Load and process the EUNIS to Red List crosswalk data.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        dict: A dictionary mapping EUNIS codes to Red List codes.
    """
    eunis_habitats_filepath = args.datasets_filepath + 'eunis_habitats.xlsx'  # Construct the filepath of the EUNIS habitats Excel file
    eunis_habitats = pd.read_excel(eunis_habitats_filepath)  # Read the Excel file into a pandas DataFrame
    eunis_red_list_crosswalks = dict(zip(eunis_habitats['EUNIS 2020 code'], eunis_habitats['Red List code']))  # Create a dictionary mapping EUNIS codes to Red List codes
    eunis_red_list_crosswalks = {key: eunis_red_list_crosswalks[key] for key in eunis_red_list_crosswalks if type(eunis_red_list_crosswalks[key]) is str}  # Filter the crosswalk dictionary to keep only string values
    return eunis_red_list_crosswalks

def get_endangered_red_list_habitats(args):
    """
    Load and process the endangered Red List habitats data.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        dict: A dictionary mapping Red List project habitat codes to endangered habitat categories.
    """
    red_list_habitats_filepath = args.datasets_filepath + 'red_list_habitats.xlsx'  # Construct the filepath of the Red List habitats Excel file
    red_list_habitats = pd.read_excel(red_list_habitats_filepath, sheet_name='Terrestrial')  # Read the Excel file into a pandas DataFrame
    red_list_habitats = dict(zip(red_list_habitats['RL project habitat code'], red_list_habitats['Overall category EU28']))  # Create a dictionary mapping Red List project habitat codes to overall categories
    endangered_red_list_habitats = {key: val for key, val in red_list_habitats.items() if 'Endangered' in val}  # Filter the habitats dictionary to keep only endangered categories
    return endangered_red_list_habitats

def get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header):
    """
    Get the encoded labels of endangered EUNIS habitats based on crosswalks and label encoder.

    Args:
        dict: A dictionary mapping EUNIS codes to Red List codes.
        dict: A dictionary mapping Red List project habitat codes to endangered habitat categories.
        sklearn.preprocessing._label.LabelEncoder: The label encoder object for encoding habitat labels.

    Returns:
        list: Encoded labels of endangered EUNIS habitats.
    """
    endangered_habitats_with_crosswalks = {key: value for key, value in endangered_red_list_habitats.items() if key in eunis_red_list_crosswalks.values()}  # Filter the endangered Red List habitats to keep only those with crosswalks
    endangered_eunis_habitats = [key for key in eunis_red_list_crosswalks.keys() if eunis_red_list_crosswalks[key] in endangered_habitats_with_crosswalks.keys()]  # Find the EUNIS codes of the endangered habitats
    endangered_eunis_habitats = [key for key in endangered_eunis_habitats if key in le_header.classes_]  # Filter the EUNIS codes to keep only those present in the label encoder classes
    endangered_eunis_habitats = le_header.transform(endangered_eunis_habitats)  # Encode the EUNIS habitat labels
    return endangered_eunis_habitats

def get_ohe_country(args):
    """
    Load the one-hot encoding representation of countries.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        sklearn.preprocessing._encoders.OneHotEncoder: The one-hot encoding representation of countries.
    """
    ohe_country_filepath = args.data_filepath + 'ohe_country.pkl'  # Filepath of the one-hot encoding pickle file
    with open(ohe_country_filepath, 'rb') as f:  # Open the pickle file in binary mode for reading
        ohe_country = pickle.load(f)  # Load the contents of the pickle file
    return ohe_country

def get_ohe_ecoregion(args):
    """
    Load the one-hot encoding representation of ecoregions.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        sklearn.preprocessing._encoders.OneHotEncoder: The one-hot encoding representation of ecoregions.
    """
    ohe_ecoregion_filepath = args.data_filepath + 'ohe_ecoregion.pkl'  # Filepath of the one-hot encoding pickle file
    with open(ohe_ecoregion_filepath, 'rb') as f:  # Open the pickle file in binary mode for reading
        ohe_ecoregion = pickle.load(f)  # Load the contents of the pickle file
    return ohe_ecoregion

def get_ohe_dune(args):
    """
    Load the one-hot encoding representation of dunes.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        sklearn.preprocessing._encoders.OneHotEncoder: The one-hot encoding representation of dunes.
    """
    ohe_dune_filepath = args.data_filepath + 'ohe_dune.pkl'  # Filepath of the one-hot encoding pickle file
    with open(ohe_dune_filepath, 'rb') as f:  # Open the pickle file in binary mode for reading
        ohe_dune = pickle.load(f)  # Load the contents of the pickle file
    return ohe_dune

def get_ohe_coast(args):
    """
    Load the one-hot encoding representation of coast categories.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        sklearn.preprocessing._encoders.OneHotEncoder: The one-hot encoding representation of coast categories.
    """
    ohe_coast_filepath = args.data_filepath + 'ohe_coast.pkl'  # Filepath of the one-hot encoding pickle file
    with open(ohe_coast_filepath, 'rb') as f:  # Open the pickle file in binary mode for reading
        ohe_coast = pickle.load(f)  # Load the contents of the pickle file
    return ohe_coast

def get_scaler(args):
    """
    Load the scaler object used for feature scaling.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.

    Returns:
        sklearn.preprocessing._data.StandardScaler: The scaler object used for feature scaling.
    """
    scaler_filepath = args.data_filepath + 'scaler.pkl'  # Filepath of the scaler pickle file
    with open(scaler_filepath, 'rb') as f:  # Open the pickle file in binary mode for reading
        scaler = pickle.load(f)  # Load the contents of the pickle file
    return scaler

def get_mlp_model(args, model):
    """
    Load the pre-trained MLP model.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.
        model (models.mlp.MLP | torch.nn.parallel.data_parallel.DataParallel): The MLP model object.

    Returns:
        models.mlp.MLP | torch.nn.parallel.data_parallel.DataParallel: The loaded pre-trained MLP model.
    """
    mlp_filepath = args.models_filepath + 'MLP.pth'  # Filepath of the pre-trained MLP model
    model.load_state_dict(torch.load(mlp_filepath))  # Load the weights of the pre-trained model
    return model

def get_rfc_model(args, model):
    """
    Load the pre-trained Random Forest Classifier (RFC) model.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.
        models.rfc.RFC: The RFC model object.

    Returns:
        models.rfc.RFC: The loaded pre-trained RFC model.
    """
    rfc_filepath = args.models_filepath + 'RFC.joblib'  # Filepath of the pre-trained RFC model
    model = joblib.load(rfc_filepath)  # Load the pre-trained RFC model
    return model

def get_xgb_model(args, model):
    """
    Load the pre-trained XGBoost model.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.
        models.xgb.XGB: The XGBoost model object.

    Returns:
        models.xgb.XGB: The loaded pre-trained XGBoost model.
    """
    xgb_filepath = args.models_filepath + 'XGB.json'  # Filepath of the pre-trained XGBoost model
    model.load_model(xgb_filepath)  # Load the pre-trained XGBoost model
    return model

def get_tnc_model(args, model):
    """
    Load the pre-trained TNC model.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.
        models.tnc.TNC: The TNC model object.

    Returns:
        models.tnc.TNC: The loaded pre-trained TNC model.
    """
    tnc_filepath = args.models_filepath + 'TNC.zip'  # Construct the filepath of the TNC model
    model.load_model(tnc_filepath)  # Load the pre-trained TNC model from the file path
    return model

def get_ftt_model(args, model):
    """
    Load the pre-trained FTT model.

    Args:
        args (argparse.Namespace): An object containing the necessary arguments.
        models.ftt.FTT: The FTT model object.

    Returns:
        models.ftt.FTT: The loaded pre-trained FTT model.
    """
    ftt_filepath = args.models_filepath + 'FTT.pth'  # Construct the filepath of the FTT model
    model.load_state_dict(torch.load(ftt_filepath))  # Load the pre-trained FTT model from the file path
    return model

data/load_data.py
cesar-leblanc/hdm-framework
import numpy as np
import sklearn.preprocessing
import random
import imblearn
import collections
import pandas as pd
import requests
import tqdm
import scipy.sparse
import rtree
import shapely
import geopandas as gpd
import torch
import rioxarray
import pyproj

from utils import assign_block_ids, SparseDataset, get_available_subprocesses

from data.load_data import get_le_species, get_ohe_country, get_ohe_ecoregion, get_ohe_dune, get_ohe_coast

def add_fold_assignments(X, y, split_assignments, fold):
    """
    Split the data into training and testing sets based on fold assignments.

    Args:
        X (scipy.sparse._csr.csr_matrix): The input features.
        y (numpy.ndarray): The target labels.
        split_assignments (numpy.ndarray): The fold assignments for each data point.
        fold (int): The fold number to use as the testing set.

    Returns:
        tuple: A tuple containing the training and testing sets for both X and y.
    """
    X_train = X[split_assignments != fold]  # Select data points where split_assignments is not equal to the given fold
    X_test = X[split_assignments == fold]  # Select data points where split_assignments is equal to the given fold
    y_train = y[split_assignments != fold]  # Select corresponding target labels for the training set
    y_test = y[split_assignments == fold]  # Select corresponding target labels for the testing set
    return X_train, X_test, y_train, y_test

def add_dataloader(args, X, y, shuffle=True):
    """
    Create a DataLoader for the given input features and target labels.

    Args:
        args (argparse.Namespace): Command-line arguments.
        X (scipy.sparse._csr.csr_matrix): The input features.
        y (numpy.ndarray): The target labels.
        shuffle (bool, optional): Whether to shuffle the data. Defaults to True.

    Returns:
        torch.utils.data.DataLoader: The DataLoader object for iterating over the data in batches.
    """
    dataset = SparseDataset(X, y)  # Create a dataset from the input features and target labels
    dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=args.batch_size, shuffle=shuffle, num_workers=get_available_subprocesses(args))  # Create a DataLoader for the dataset with the specified batch size, shuffle, and number of workers
    return dataloader

def add_standardization(X, scaler=None):
    """
    Apply standardization to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): The input data.
        scaler (sklearn.preprocessing._data.StandardScaler, optional): The scaler object. If None, a new StandardScaler instance will be created and fitted on the data. Defaults to None.

    Returns:
        tuple: A tuple containing the standardized data (X) and the scaler object.
    """
    if not scaler:
        scaler = sklearn.preprocessing.StandardScaler(with_mean=False)  # Create a StandardScaler instance
        scaler.fit(X)  # Compute the mean and std on the training data to be used for later scaling
    X = scaler.transform(X)  # Scale the data using the scaler
    return X, scaler

def add_normalization(X, n_columns):
    """
    Apply normalization to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): The input data in CSR sparse matrix format.
        n_columns (int): The number of columns in the input data.

    Returns:
        scipy.sparse._csr.csr_matrix: The normalized input data in CSR sparse matrix format.
    """
    normalizer = sklearn.preprocessing.Normalizer(norm='l1')  # Create a Normalizer instance with the l1 norm
    for i in range(X.shape[0]):  # Iterate over the rows
        start, end = X.indptr[i], X.indptr[i+1]  # The indptr attribute of the CSR matrix defines the start and end indices of each row
        indices = X.indices[start:end]  # For each row, retrieve the indices of the non-zero elements using the 'indices' attribute of the CSR matrix
        data = X.data[start:end]  # Retrieve the non-zero elements of the row using the 'data' attribute of the CSR matrix
        mask = indices < n_columns  # Create a mask for the selected columns
        selected_data = data[mask]  # Use the mask to retrieve the non-zero elements of the row that correspond to the columns we want
        normalizer.fit(selected_data.reshape(1, -1))  # Do nothing and return the estimator unchanged
        selected_data = normalizer.transform(selected_data.reshape(1, -1))  # Normalize the matrix (scale each non zero row to unit norm)
        data[mask] = selected_data[0]  # Update the non-zero elements of the row in the CSR matrix
    return X

def add_dropout(X, p, n_columns):
    """
    Apply dropout to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): The input data in CSR sparse matrix format.
        p (float): The dropout probability, indicating the proportion of non-zero elements to be set to zero.
        n_columns (int): The number of columns for species in the input data.

    Returns:
        scipy.sparse._csr.csr_matrix: The input data with dropout applied in CSR sparse matrix format.
    """
    column_indices = np.where(np.isin(X.indices, list(range(n_columns))))[0]  # Get the indices of the non-zero elements in the columns that we want
    sample_size = int(len(column_indices) * p)  # Set the sample size to p% of the total number of non-zero elements
    sample_indices = random.sample(list(column_indices), sample_size)  # Select a random subset of these indices
    X.data[sample_indices] = 0  # Set the values at the selected indices to 0 
    X.eliminate_zeros()  # Remove explicit zero elements from the matrix
    return X

def add_binarization(X, n_columns):
    """
    Apply binarization to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): The input data in CSR sparse matrix format.
        n_columns (int): The number of columns for species in the input data.

    Returns:
        scipy.sparse._csr.csr_matrix: The input data with binarization applied in CSR sparse matrix format.
    """
    binarizer = sklearn.preprocessing.Binarizer(threshold=1)  # Create a binarizer instance with threshold=1
    column_indices = list(range(n_columns))  # Select the columns we want to binarize
    X.data[np.isin(X.indices, column_indices)] = binarizer.transform(X.data[np.isin(X.indices, column_indices)].reshape(-1, 1)).flatten()  # Binarize the data in the selected columns
    X.eliminate_zeros()  # Remove explicit zero elements from the matrix
    return X

def add_log(X, n_columns):
    """
    Apply natural logarithm to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): The input data in CSR sparse matrix format.
        n_columns (int): The number of columns for species in the input data.

    Returns:
        scipy.sparse._csr.csr_matrix: The input data with natural logarithm applied to the selected columns in CSR sparse matrix format.
    """
    column_indices = list(range(n_columns))  # Select the columns we want to apply the natural logarithm to
    mask = np.isin(X.indices, column_indices)  # Create a mask for the selected columns
    X.data[mask] = np.log(X.data[mask])  # Use the mask to select the columns and apply the natural logarithm
    return X

def add_noise(X, n_columns):
    """
    Add noise to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): The input data in CSR sparse matrix format.
        n_columns (int): The number of columns for species in the input data.

    Returns:
        scipy.sparse._csr.csr_matrix: The input data with noise added to the selected columns in CSR sparse matrix format.
    """
    mean_column = np.zeros(n_columns)  # Initialize mean for each column
    std_column = np.zeros(n_columns)  # Initialize standard deviation for each column
    for j in range(n_columns):  # Iterate through the columns we will modify
        mean_column[j] = np.mean(X[:, j].data)  # Calculate and store the mean of the elements in the current column
        std_column[j] = np.std(X[:, j].data)  # Calculate and store the standard deviation of the elements in the current column
    for i in range(len(X.data)):  # Iterate through the non-zero elements of X
        if X.indices[i] < n_columns:  # Check if the element is in one of the columns we want to modify
            noise = np.random.normal(mean_column[X.indices[i]], std_column[X.indices[i]])  # Generate noise value from normal distribution with mean and standard deviation calculated from the same column
            noise *= np.random.choice([-1, 1])  # Randomly choose whether to add or subtract the noise
            X.data[i] += noise  # Add the noise value to the element
            X.data[i] = np.max(X.data[i], 0)  # Ensure that the element is non-negative (clip the value at zero)
    X.eliminate_zeros()  # Eliminate explicit zeros from the sparse matrix to save memory
    return X

def add_rank(X, n_columns):
    """
    Add rank-based transformation to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): The input data in CSR sparse matrix format.
        n_columns (int): The number of columns for species in the input data.

    Returns:
        scipy.sparse._csr.csr_matrix: The input data with rank-based transformation applied to the selected columns in CSR sparse matrix format.
    """
    for i in range(X.shape[0]):  # Iterate over the rows
        start, end = X.indptr[i], X.indptr[i+1]  # The indptr attribute of the CSR matrix defines the start and end indices of each row
        indices = X.indices[start:end]  # For each row, retrieve the indices of the non-zero elements using the 'indices' attribute of the CSR matrix
        data = X.data[start:end]  # Retrieve the non-zero elements of the row using the 'data' attribute of the CSR matrix
        mask = indices < n_columns  # Create a mask for the selected columns
        selected_data = data[mask]  # Use the mask to retrieve the non-zero elements of the row that correspond to the columns we want
        sorted_data = sorted(np.unique(selected_data), reverse=True)  # Sort the unique elements of the array in descending order
        rank_dict = {val: rank for rank, val in enumerate(sorted_data, start=1)}  # Use the sorted array to create a dictionary that maps each value to its rank
        selected_data = [1/rank_dict[val] for val in selected_data]  # Transform the original array by replacing each value with its corresponding rank from the dictionary, and dividing 1 by the rank
        data[mask] = selected_data  # Update the non-zero elements of the row in the CSR matrix
    return X

def remove_features(args, X, features):
    """
    Remove specific features from the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): The input data in CSR sparse matrix format.
        features (list): A list of features to remove from the data.

    Returns:
        scipy.sparse._csr.csr_matrix: The input data with the specified features removed in CSR sparse matrix format.
    """
    le_species = get_le_species(args)  # Get the LabelEncoder for species
    ohe_country = get_ohe_country(args)  # Get the OneHotEncoder for country
    ohe_ecoregion = get_ohe_ecoregion(args)  # Get the OneHotEncoder for ecoregion
    ohe_dune = get_ohe_dune(args)  # Get the OneHotEncoder for dune
    ohe_coast = get_ohe_coast(args)  # Get the OneHotEncoder for coast
    
    select_columns = []  # Initialize a list to store the selected columns
    
    if 'species' in features:
        select_columns += [*range(0, len(le_species.classes_), 1)]  # Add the indices of species columns to the selected columns
    if 'location' in features:
        select_columns += [*range(len(le_species.classes_), len(le_species.classes_) + 2, 1)]  # Add the indices of location columns to the selected columns    
    if 'altitude' in features:
        select_columns += [*range(len(le_species.classes_) + 2, len(le_species.classes_) + 2 + 1, 1)]  # Add the index of altitude column to the selected columns
    if 'country' in features:
        select_columns += [*range(len(le_species.classes_) + 2 + 1, len(le_species.classes_) + 2 + 1 + len(ohe_country.categories_[0]), 1)]  # Add the indices of country columns to the selected columns
    if 'ecoregion' in features:
        select_columns += [*range(len(le_species.classes_) + 2 + 1 + len(ohe_country.categories_[0]), len(le_species.classes_) + 2 + 1 + len(ohe_country.categories_[0]) + len(ohe_ecoregion.categories_[0]), 1)]  # Add the indices of ecoregion columns to the selected columns
    if 'dune' in features:
        select_columns += [*range(len(le_species.classes_) + 2 + 1 + len(ohe_country.categories_[0]) + len(ohe_ecoregion.categories_[0]), len(le_species.classes_) + 2 + 1 + len(ohe_country.categories_[0]) + len(ohe_ecoregion.categories_[0]) + len(ohe_dune.categories_[0]), 1)]  # Add the indices of dune columns to the selected columns
    if 'coast' in features:
        select_columns += [*range(len(le_species.classes_) + 2 + 1 + len(ohe_country.categories_[0]) + len(ohe_ecoregion.categories_[0]) + len(ohe_dune.categories_[0]), len(le_species.classes_) + 2 + 1 + len(ohe_country.categories_[0]) + len(ohe_ecoregion.categories_[0]) + len(ohe_dune.categories_[0]) + len(ohe_coast.categories_[0]), 1)]  # Add the indices of coast columns to the selected columns
    if 'all' in features:
        select_columns = [*range(0, len(le_species.classes_) + 2 + 1 + len(ohe_country.categories_[0]) + len(ohe_ecoregion.categories_[0]) + len(ohe_dune.categories_[0]) + len(ohe_coast.categories_[0]), 1)]

    X = X[:,select_columns]  # Select the columns specified in select_columns from the input data

    return X

def add_augmentation(X, y, labels=None):
    """
    Apply augmentation to the input data by oversampling minority classes using BorderlineSMOTE.

    Args:
        X (scipy.sparse._csr.csr_matrix): The input data.
        y (numpy.ndarray): The target labels.
        labels (list, optional): The specific labels to augment. If None, all minority classes will be augmented.

    Returns:
        scipy.sparse._csr.csr_matrix: The augmented input data.
        numpy.ndarray: The augmented target labels.
    """
    n_samples_majority = max(np.bincount(y))  # Calculate the number of samples in the majority class
    n_samples_minority = int(n_samples_majority * 0.1)  # Set the target number of samples for the minority classes to 10% of the number of samples in the majority class
    class_counts = dict(collections.Counter(y))  # Get the number of samples for each class in the original data
    if labels is None:
        sampling_strategy = {i: max(n, n_samples_minority) for i, n in class_counts.items()}  # Define the sampling strategy as a dictionary where each class has at least the maximum of its original number of samples and n_samples_minority
    else:
        sampling_strategy = {i: max(n, n_samples_minority) for i, n in class_counts.items() if i in labels}
    k_neighbors = int(min(class_counts.values()) / 2)  # Define the number of neighbors to be considered for each sample
    m_neighbors = k_neighbors * 2  # Define the number of neighbors to consider to determine if a sample is danger
    sm = imblearn.over_sampling.BorderlineSMOTE(sampling_strategy=sampling_strategy, k_neighbors=k_neighbors, m_neighbors=m_neighbors)  # Define the oversampler
    X_resampled, y_resampled = sm.fit_resample(X, y)  # Fit the oversampler on the training data
    return X_resampled, y_resampled 

def add_rare_labels(y, X_test, y_test):
    """
    Filter out rare labels from the test data based on their occurrence in the training labels.

    Args:
        y (numpy.ndarray): The training labels.
        X_test (scipy.sparse._csr.csr_matrix): The test data.
        y_test (numpy.ndarray): The test labels.

    Returns:
        scipy.sparse._csr.csr_matrix: The filtered test data containing samples with rare labels.
        numpy.ndarray: The filtered test labels containing rare labels.
    """
    labels, counts = np.unique(y, return_counts=True)  # Get the unique labels and their respective counts in the training labels
    sorted_labels = labels[np.argsort(counts)]  # Sort the labels based on their counts in ascending order
    rare_labels = sorted_labels[:10]  # Select the 10 labels with the lowest counts (rare labels)
    rare_indices = np.isin(y_test, rare_labels)  # Create a boolean mask indicating whether each label in the test labels is a rare label
    X_test_rare = X_test[rare_indices]  # Filter the test data based on the rare label indices
    y_test_rare = y_test[rare_indices]  # Filter the test labels based on the rare label indices
    return X_test_rare, y_test_rare

def add_endangered_habitats(X, y, endangered_eunis_habitats):
    """
    Filter the data and labels to include only samples from endangered habitats.

    Args:
        X (scipy.sparse._csr.csr_matrix): The data matrix.
        y (numpy.ndarray): The labels.
        endangered_eunis_habitats (list): The list of endangered EUNIS habitats.

    Returns:
        scipy.sparse._csr.csr_matrix: The filtered data matrix containing samples from endangered habitats.
        numpy.ndarray: The filtered labels corresponding to the filtered data.
    """
    mask = np.isin(y, endangered_eunis_habitats)  # Create a Boolean mask indicating which samples belong to endangered habitats
    X = X[mask]  # Use the mask to extract only the desired samples from the X matrix
    y = y[mask]  # Use the mask to extract only the desired labels from the y array
    return X, y

def add_header_encoding(eva_header):
    """
    Encode the 'Expert System' column in the EVA header using LabelEncoder.

    Args:
        eva_header (pandas.core.frame.DataFrame): The EVA header DataFrame.

    Returns:
        pandas.core.frame.DataFrame: The modified EVA header DataFrame with the 'Expert System' column encoded.
        sklearn.preprocessing._label.LabelEncoder: The LabelEncoder object used for encoding.
    """
    le_header = sklearn.preprocessing.LabelEncoder()  # Create a LabelEncoder object
    le_header.fit(eva_header['Expert System'])  # Fit the LabelEncoder on the 'Expert System' column
    eva_header['Expert System'] = le_header.transform(eva_header['Expert System'])  # Encode the 'Expert System' column
    return eva_header, le_header

def add_header_decoding(eva_header, le_header):
    """
    Decode the encoded 'Expert System' column in the EVA header using a LabelEncoder.

    Args:
        eva_header (pandas.core.frame.DataFrame): The EVA header DataFrame.
        le_header (sklearn.preprocessing._label.LabelEncoder): The LabelEncoder object used for encoding.

    Returns:
        pandas.core.frame.DataFrame: The modified EVA header DataFrame with the 'Expert System' column decoded.
    """
    eva_header['Expert System'] = le_header.inverse_transform(eva_header['Expert System'])  # Decode the 'Expert System' column using the inverse_transform method of the LabelEncoder
    return eva_header

def add_predictions_decoding(predictions, le_header):
    """
    Decode the encoded predictions using a LabelEncoder.

    Args:
        predictions (list): The encoded predictions.
        le_header (sklearn.preprocessing._label.LabelEncoder): The LabelEncoder object used for encoding.

    Returns:
        numpy.ndarray: The decoded predictions.

    """
    predictions = le_header.inverse_transform(predictions)  # Decode the predictions using the inverse_transform method of the LabelEncoder
    return predictions

def add_gbif_normalization(df_species, eva_to_gbif_species):
    """
    Perform GBIF species name normalization on a DataFrame column.

    Args:
        df_species (pandas.core.frame.DataFrame): DataFrame containing the species data.
        eva_to_gbif_species (dict): Mapping of EVA species names to corresponding GBIF species names.

    Returns:
        pandas.core.frame.DataFrame: DataFrame with normalized species names.
        dict: Mapping of original species names to normalized GBIF species names.
    """
    base = "https://api.gbif.org/v1"  # Define the base URL
    api = "species"  # Define the GBIF API we want to query
    function = "match"  # Define the functionality we want to use
    parameter = "name"  # Define the parameters for our API call
    url = f"{base}/{api}/{function}?{parameter}="  # Define the URL
    species_df = pd.unique(df_species['Matched concept'])  # Get unique species names from the DataFrame column
    species_gbif = []  # Initilize an empty list to store GBIF species names
    print()
    for species in tqdm.tqdm(species_df, desc="GBIF normalization"):  # Iterate over species names, displaying progress with tqdm
        url = url.replace(url.partition('name')[2], f'={species}')  # Construct the URL for GBIF API with the species name
        r = requests.get(url)  # Send a GET request to GBIF API
        r = r.json()  # Parse the response as JSON
        if 'phylum' in r and r['phylum'] == 'Tracheophyta' and 'species' in r:
            r = r["species"]  # Retrieve the GBIF species name
            if eva_to_gbif_species and r not in eva_to_gbif_species.values():
                r = '?'  # If eva_to_gbif_species is provided and the species is not in the mapping, replace with '?'
        else:
            r = '?'  # If the species is not found in GBIF, replace with '?'
        species_gbif.append(r)  # Append the normalized species name to the list
    df_to_gbif_species = dict(zip(species_df, species_gbif))  # Create a mapping of original species names to normalized GBIF species names
    df_species['Matched concept'] = df_species['Matched concept'].map(df_to_gbif_species)  # Map the normalized species names to the DataFrame column
    df_species = df_species[df_species["Matched concept"] != '?']  # Remove rows with '?' as the normalized species name
    df_species = df_species.groupby(['PlotObservationID', 'Matched concept']).sum().reset_index()  # Group by plot observation ID and normalized species name, summing other columns
    return df_species, df_to_gbif_species

def add_species_encoding(df_species, le_species):
    """
    Encode species names in a DataFrame column using LabelEncoder.

    Args:
        df_species (pandas.core.frame.DataFrame): DataFrame containing the species data.
        le_species (sklearn.preprocessing._label.LabelEncoder, optional): Pre-initialized LabelEncoder object for species encoding.

    Returns:
        pandas.core.frame.DataFrame: DataFrame with encoded species names.
        sklearn.preprocessing._label.LabelEncoder: LabelEncoder object used for encoding.
    """
    if not le_species:
        le_species = sklearn.preprocessing.LabelEncoder()  # If LabelEncoder object is not provided, initialize a new one
        le_species.fit(df_species['Matched concept'])  # Fit the LabelEncoder on the species data
    df_species['Matched concept'] = le_species.transform(df_species['Matched concept'])  # Encode the species names using the LabelEncoder
    return df_species, le_species

def add_species_decoding(df_species, le_species):
    """
    Decode encoded species names in a DataFrame column using a given LabelEncoder.

    Args:
        df_species (pandas.core.frame.DataFrame): DataFrame containing the species data.
        le_species (sklearn.preprocessing._label.LabelEncoder): LabelEncoder object used for encoding species names.

    Returns:
        pandas.core.frame.DataFrame: DataFrame with decoded species names.
    """
    df_species['Matched concept'] = le_species.inverse_transform(df_species['Matched concept'])  # Decode the encoded species names using the inverse_transform method of the LabelEncoder
    return df_species

def add_spatial_split(args, eva_header):
    """
    Assign spatial blocks and split data into folds based on spatial coordinates.

    Args:
        args (argparse.Namespace): Command-line arguments or configuration parameters.
        eva_header (pandas.core.frame.DataFrame): DataFrame containing the habitat data.

    Returns:
        numpy.ndarray: Array representing the fold assignments for each data point.
        pandas.core.frame.DataFrame: DataFrame with an additional column for fold assignments.
    """
    lon_list = eva_header['Longitude'].to_numpy()
    lat_list = eva_header['Latitude'].to_numpy()
    east_west_bin_width_km = args.east_west_bin_width_km
    north_south_bin_width_km = args.north_south_bin_width_km
    origin = (0.0, 0.0)  # Define the origin
    block_id_assignments = assign_block_ids(lon_list, lat_list, east_west_bin_width_km, north_south_bin_width_km, origin)  # Assign spatial block IDs to each data point based on coordinates
    vals = np.unique(block_id_assignments)
    print(f"There are {len(vals)} different blocks.")
    block_ids = np.unique(block_id_assignments)
    folds = [len(block_ids) // args.folds + (1 if x < len(block_ids) % args.folds else 0) for x in range(args.folds)]  # Split the block IDs into folds
    idx_rand = np.random.permutation(len(block_ids))  # Randomly permute the block IDs
    idx = np.array_split(idx_rand, args.folds)  # Split the permuted block IDs into folds
    assert len(np.unique(np.concatenate(idx).ravel())) == len(block_ids)  # Ensure that the concatenated indices cover all block IDs
    folds_blocks_id = [block_ids[idx[i]] for i in range(len(idx))]  # Create a list of block IDs for each fold
    assert len(np.concatenate(folds_blocks_id)) == len(block_ids)  # Ensure that the concatenated block IDs cover all unique block IDs
    split_assignment = [fold for i in range(len(lon_list)) for fold in range(args.folds) if block_id_assignments[i] in folds_blocks_id[fold]]  # Assign fold numbers to each data point based on the block IDs
    split_assignment = np.array(split_assignment, dtype=int)
    print(f'{args.folds} different folds:')
    for i in range(args.folds):
        print(f'Fold {i}: {folds[i]} blocks')
    eva_header['Fold'] = split_assignment  # Add a column for fold assignments to the DataFrame
    return split_assignment, eva_header

def add_input_data(df_header, df_species, le_species):
    """
    Create input data matrix for habitat distribution modeling.

    Args:
        df_header (pandas.core.frame.DataFrame): DataFrame containing the header data.
        df_species (pandas.core.frame.DataFrame): DataFrame containing the species data.
        le_species (sklearn.preprocessing._label.LabelEncoder): LabelEncoder object for species encoding.

    Returns:
        scipy.sparse._csr.csr_matrix: Input data matrix with species coverage information.
    """
    data = []  # Initialize the data list
    indices = []  # Initialize the indices list
    X = df_header["PlotObservationID"].values
    for i in tqdm.tqdm(range(len(X)), desc="Input data"):
        species_cover = df_species[df_species["PlotObservationID"] == X[i]].values[:, 2]  # Extract species coverage for a specific observation ID
        species_ind = df_species[df_species["PlotObservationID"] == X[i]].values[:, 1]  # Extract species indices for a specific observation ID
        data.append(species_cover)  # Append species coverage to the data list
        indices.append(species_ind)  # Append species indices to the indices list
    indptr = np.empty(len(df_header) + 1)  # Create indptr array for efficient sparse matrix construction
    indptr[0] = 0
    for i in range(1, len(indptr)):
        indptr[i] = indptr[i-1] + len(data[i-1])
    data = np.concatenate(data)  # Concatenate the data array
    indices = np.concatenate(indices)  # Concatenate the indices array
    X = scipy.sparse.csr_matrix((data, indices, indptr))  # Create a compressed sparse row (CSR) matrix using the data, indices, and indptr arrays
    n = len(le_species.classes_)  # Desired number of columns
    num_rows, num_cols = X.shape
    if num_cols < n:  # Check if the number of columns in X is less than n
        num_zeros = n - num_cols
        zeros = scipy.sparse.csr_matrix((num_rows, num_zeros), dtype=X.dtype)
        X = scipy.sparse.hstack([X, zeros])
    assert X.shape == (len(df_header), len(le_species.classes_))  # Ensure the shape of X matches the expected shape
    return X

def add_target_values(eva_header):
    """
    Extract target values for habitat distribution modeling.

    Args:
        eva_header (pandas.core.frame.DataFrame): DataFrame containing the header data.

    Returns:
        numpy.ndarray: Array of target values.
    """
    y = eva_header['Expert System'].values  # Extract the target values by accessing the 'Expert System' column
    return y

def add_longitude(X, df_header):
    """
    Add longitude feature to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): Input data matrix.
        df_header (pandas.core.frame.DataFrame): DataFrame containing the header data.

    Returns:
        scipy.sparse._csr.csr_matrix: Updated input data matrix with longitude feature.
    """
    X = scipy.sparse.hstack((X, np.array(df_header['Longitude'])[:,None]), format="csr")  # Concatenate the longitude array with the input data matrix horizontally
    return X

def add_latitude(X, df_header):
    """
    Add latitude feature to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): Input data matrix.
        df_header (pandas.core.frame.DataFrame): DataFrame containing the header data.

    Returns:
        scipy.sparse._csr.csr_matrix: Updated input data matrix with latitude feature.
    """
    X = scipy.sparse.hstack((X, np.array(df_header['Latitude'])[:,None]), format="csr")  # Concatenate the latitude array with the input data matrix horizontally
    return X

def add_altitude(X, df_header, altitudes):
    """
    Add altitude feature to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): Input data matrix.
        df_header (pandas.core.frame.DataFrame): DataFrame containing the header data.
        altitudes (xarray.core.dataarray.DataArray): DataArray containing altitude values.

    Returns:
        scipy.sparse._csr.csr_matrix: Updated input data matrix with altitude feature.
        pandas.core.frame.DataFrame: Updated DataFrame with altitude values.
    """
    transformer = pyproj.Transformer.from_crs("EPSG:4326", "EPSG:3857")  # Create a coordinate transformer from EPSG:4326 (latitude and longitude) to EPSG:3857 (Web Mercator)
    list_of_altitudes = []  # Initialize a list to store the altitude values
    for i in tqdm.tqdm(range(len(df_header)), desc="Altitudes"):  # Iterate over each row in the df_header DataFrame
        x, y = transformer.transform(df_header['Latitude'][i], df_header['Longitude'][i])  # Transform latitude and longitude to Web Mercator coordinates
        list_of_altitudes.append(altitudes.sel(x=x, y=y, method='nearest').item())  # Retrieve the nearest altitude value from the altitudes DataArray based on the transformed coordinates and append it to the list
    df_header['Altitude'] = list_of_altitudes   # Add the altitude values as a new column 'Altitude' to the df_header DataFrame
    list_of_altitudes = scipy.sparse.csr_matrix(list_of_altitudes).T  # Convert the list of altitudes to a sparse matrix with a single column (transpose)
    X = scipy.sparse.hstack((X, list_of_altitudes), format='csr')  # Horizontally stack the altitude sparse matrix with the input data matrix X
    return X, df_header

def add_country(args, X, y, df_header, df_species, countries, ohe_country):
    """
    Add country feature to the input data.

    Args:
        args (argparse.Namespace): Arguments for controlling country filtering.
        X (scipy.sparse._csr.csr_matrix): Input data matrix.
        y (numpy.ndarray): Array of target values.
        df_header (pandas.core.frame.DataFrame): DataFrame containing the header data.
        df_species (pandas.core.frame.DataFrame): DataFrame containing the species data.
        countries (geopandas.geodataframe.GeoDataFrame): GeoDataFrame containing country boundaries.
        ohe_country (sklearn.preprocessing._encoders.OneHotEncoder): OneHotEncoder for country labels.

    Returns:
        scipy.sparse._csr.csr_matrix: Updated input data matrix with country feature.
        numpy.ndarray: Array of target values.
        sklearn.preprocessing._encoders.OneHotEncoder: Updated OneHotEncoder for country labels.
        pandas.core.frame.DataFrame: Updated header DataFrame with country labels.
        pandas.core.frame.DataFrame: Updated species DataFrame with observations from specific countries.
    """
    country_index = rtree.index.Index()  # Create a spatial index for the country boundaries
    for index, country in countries.iterrows():  # Insert each country's index and bounding box into the spatial index
        country_index.insert(index, country["geometry"].bounds)
    list_of_countries = []   # Initialize a list to store the country labels
    for i in tqdm.tqdm(range(len(df_header)), desc="Countries"):  # Iterate over each row in the df_header DataFrame
        point = shapely.geometry.Point(df_header['Longitude'][i], df_header['Latitude'][i])
        point = {'geometry': [point]}  # Create a point geometry from the longitude and latitude
        point = gpd.GeoDataFrame(point, crs="EPSG:4326")  # Convert the point to a GeoDataFrame with EPSG:4326 CRS
        point = point.to_crs(crs=3857)  # Convert the point's CRS to EPSG:3857 (Web Mercator)
        point = pd.Series(data=point['geometry'][0], index=['geometry'])  # Extract the geometry of the point as a Pandas Series
        min_distance = float("inf")  # Initialize variables for finding the closest country
        closest_country = None
        for country_id in country_index.nearest((point["geometry"].x, point["geometry"].y, point["geometry"].x, point["geometry"].y), 1):  # Find countries that are within a certain distance of the point
            distance = point["geometry"].distance(countries.iloc[country_id]["geometry"])  # Calculate the distance between the point and the country
            if distance < min_distance:  # Update the closest country if the distance is smaller
                min_distance = distance
                closest_country = countries.iloc[country_id]
        closest_country = closest_country[0]  # Get the name of the closest country
        list_of_countries.append(closest_country)  # Append the closest country to the list of countries
    if not ohe_country:  # If not using OneHotEncoder for country labels
        countries_to_keep = args.countries
        if countries_to_keep == ['all']:  # If all countries are to be kept
            countries_to_keep = ['Afghanistan', 'Albania', 'Algeria', 'American Samoa', 'Andorra', 'Angola', 'Anguilla', 'Antarctica', 'Antigua and Barbuda', 'Argentina', 'Armenia', 'Aruba', 'Australia', 'Austria', 'Azerbaijan', 'Azores', 'Bahamas', 'Bahrain', 'Bangladesh', 'Barbados', 'Belarus', 'Belgium', 'Belize', 'Benin', 'Bermuda', 'Bhutan', 'Bolivia', 'Bonaire', 'Bosnia and Herzegovina', 'Botswana', 'Bouvet Island', 'Brazil', 'British Indian Ocean Territory', 'British Virgin Islands', 'Brunei Darussalam', 'Bulgaria', 'Burkina Faso', 'Burundi', 'Cabo Verde', 'Cambodia', 'Cameroon', 'Canada', 'Canarias', 'Cayman Islands', 'Central African Republic', 'Chad', 'Chile', 'China', 'Christmas Island', 'Cocos Islands', 'Colombia', 'Comoros', 'Congo', 'Congo DRC', 'Cook Islands', 'Costa Rica', "CÃ´te d'Ivoire", 'Croatia', 'Cuba', 'Curacao', 'Cyprus', 'Czech Republic', 'Denmark', 'Djibouti', 'Dominica', 'Dominican Republic', 'Ecuador', 'Egypt', 'El Salvador', 'Equatorial Guinea', 'Eritrea', 'Estonia', 'Eswatini', 'Ethiopia', 'Falkland Islands', 'Faroe Islands', 'Fiji', 'Finland', 'France', 'French Guiana', 'French Polynesia', 'French Southern Territories', 'Gabon', 'Gambia', 'Georgia', 'Germany', 'Ghana', 'Gibraltar', 'Glorioso Islands', 'Greece', 'Greenland', 'Grenada', 'Guadeloupe', 'Guam', 'Guatemala', 'Guernsey', 'Guinea', 'Guinea-Bissau', 'Guyana', 'Haiti', 'Heard Island and McDonald Islands', 'Honduras', 'Hungary', 'Iceland', 'India', 'Indonesia', 'Iran', 'Iraq', 'Ireland', 'Isle of Man', 'Israel', 'Italy', 'Jamaica', 'Japan', 'Jersey', 'Jordan', 'Juan De Nova Island', 'Kazakhstan', 'Kenya', 'Kiribati', 'Kuwait', 'Kyrgyzstan', 'Laos', 'Latvia', 'Lebanon', 'Lesotho', 'Liberia', 'Libya', 'Liechtenstein', 'Lithuania', 'Luxembourg', 'Madagascar', 'Madeira', 'Malawi', 'Malaysia', 'Maldives', 'Mali', 'Malta', 'Marshall Islands', 'Martinique', 'Mauritania', 'Mauritius', 'Mayotte', 'Mexico', 'Micronesia', 'Moldova', 'Monaco', 'Mongolia', 'Montenegro', 'Montserrat', 'Morocco', 'Mozambique', 'Myanmar', 'Namibia', 'Nauru', 'Nepal', 'Netherlands', 'New Caledonia', 'New Zealand', 'Nicaragua', 'Niger', 'Nigeria', 'Niue', 'Norfolk Island', 'North Korea', 'North Macedonia', 'Northern Mariana Islands', 'Norway', 'Oman', 'Pakistan', 'Palau', 'Palestinian Territory', 'Panama', 'Papua New Guinea', 'Paraguay', 'Peru', 'Philippines', 'Pitcairn', 'Poland', 'Portugal', 'Puerto Rico', 'Qatar', 'RÃ©union', 'Romania', 'Russian Federation', 'Rwanda', 'Saba', 'Saint Barthelemy', 'Saint Eustatius', 'Saint Helena', 'Saint Kitts and Nevis', 'Saint Lucia', 'Saint Martin', 'Saint Pierre and Miquelon', 'Saint Vincent and the Grenadines', 'Samoa', 'San Marino', 'Sao Tome and Principe', 'Saudi Arabia', 'Senegal', 'Serbia', 'Seychelles', 'Sierra Leone', 'Singapore', 'Sint Maarten', 'Slovakia', 'Slovenia', 'Solomon Islands', 'Somalia', 'South Africa', 'South Georgia and South Sandwich Islands', 'South Korea', 'South Sudan', 'Spain', 'Sri Lanka', 'Sudan', 'Suriname', 'Svalbard', 'Sweden', 'Switzerland', 'Syria', 'Tajikistan', 'Tanzania', 'Thailand', 'Timor-Leste', 'Togo', 'Tokelau', 'Tonga', 'Trinidad and Tobago', 'Tunisia', 'Turkiye', 'Turkmenistan', 'Turks and Caicos Islands', 'Tuvalu', 'Uganda', 'Ukraine', 'United Arab Emirates', 'United Kingdom', 'United States', 'United States Minor Outlying Islands', 'Uruguay', 'US Virgin Islands', 'Uzbekistan', 'Vanuatu', 'Vatican City', 'Venezuela', 'Vietnam', 'Wallis and Futuna', 'Yemen', 'Zambia', 'Zimbabwe']  # List of all country names
        else:
            countries_to_keep = ''.join(countries_to_keep).split(',')
        list_of_countries = ['?' if country not in countries_to_keep else country for country in list_of_countries]  # Replace non-kept countries with '?'
        mask = np.array([country != '?' for country in list_of_countries])  # Create a boolean mask for filtering the data
        X = X[mask, :]  # Apply the mask to X
        if y is not None:
            y = y[mask]  # Apply the mask to y
        df_header = df_header[mask]  # Apply the mask to df_header
        df_species = df_species[df_species['PlotObservationID'].isin(df_header['PlotObservationID'])]  # Filter the species based on vegetation plots left in the header
        df_header = df_header.reset_index(drop=True)  # Reset the index
        df_species = df_species.reset_index(drop=True)  # Reset the index
        list_of_countries = [country for country in list_of_countries if country != '?']  # Remove '?' entries from the list of countries
        ohe_country = sklearn.preprocessing.OneHotEncoder(handle_unknown='ignore')  # Initialize the OneHotEncoder for country labels
        ohe_country.fit(np.reshape(list_of_countries, (-1, 1)))  # Fit the OneHotEncoder for country labels
    df_header['Country'] = list_of_countries  # Add the country labels to df_header
    list_of_countries = ohe_country.transform(np.reshape(list_of_countries, (-1, 1)))  # Transform the country labels using OneHotEncoder
    X = scipy.sparse.hstack((X, list_of_countries))  # Concatenate the country labels with X
    return X, y, ohe_country, df_header, df_species

def add_ecoregion(X, df_header, ecoregions, ohe_ecoregion):
    """
    Adds ecoregion information to the input data.

    Args:
        X (scipy.sparse._csr.csr_matrix): Input data matrix.
        df_header (pandas.core.frame.DataFrame): DataFrame containing header information.
        ecoregions (geopandas.GeoDataFrame): GeoDataFrame containing ecoregion polygons.
        ohe_ecoregion (sklearn.preprocessing._encoders.OneHotEncoder): Pre-initialized OneHotEncoder for ecoregions.

    Returns:
        scipy.sparse._csr.csr_matrix: Updated data matrix with ecoregion information.
        sklearn.preprocessing._encoders.OneHotEncoder: Updated OneHotEncoder for ecoregions.
        pandas.core.frame.DataFrame: DataFrame with added 'Ecoregion' column.
    """
    ecoregion_index = rtree.index.Index()  # Create a spatial index for efficient querying of ecoregions
    for index, ecoregion in ecoregions.iterrows():
        ecoregion_index.insert(index, ecoregion["geometry"].bounds)
    list_of_ecoregions = []  # Initialize a list to store the closest ecoregion for each data point
    for i in tqdm.tqdm(range(len(df_header)), desc="Ecoregions"):  # Iterate over the data points and find the closest ecoregion
        point = shapely.geometry.Point(df_header['Longitude'][i], df_header['Latitude'][i])  # Create a Point object from latitude and longitude
        point = {'geometry': [point]}   # Convert the Point object to a GeoDataFrame with the same CRS as ecoregions
        point = gpd.GeoDataFrame(point, crs="EPSG:4326")
        point = point.to_crs(crs=3857)
        point = pd.Series(data=point['geometry'][0], index=['geometry'])
        min_distance = float("inf")  # Find the closest ecoregion using the spatial index
        closest_ecoregion = None
        for ecoregion_id in ecoregion_index.nearest((point["geometry"].x, point["geometry"].y, point["geometry"].x, point["geometry"].y), 1):  # Find ecoregions that are within a certain distance of the point
            distance = point["geometry"].distance(ecoregions.iloc[ecoregion_id]["geometry"])  # Calculate the distance between the point and the ecoregion
            if distance < min_distance:
                min_distance = distance
                closest_ecoregion = ecoregions.iloc[ecoregion_id]
        closest_ecoregion = closest_ecoregion[0]  # Add the closest ecoregion to the list
        list_of_ecoregions.append(closest_ecoregion)
    if not ohe_ecoregion:  # If the OneHotEncoder for ecoregions is not provided 
        ohe_ecoregion = sklearn.preprocessing.OneHotEncoder(handle_unknown='ignore')  # Initialize a new OneHotEncoder for ecoregions
        ohe_ecoregion.fit(np.reshape(list_of_ecoregions, (-1, 1)))  # Fit the new OneHotEncoder for ecoregions
    df_header['Ecoregion'] = list_of_ecoregions  # Add the ecoregion labels to df_header
    list_of_ecoregions = ohe_ecoregion.transform(np.reshape(list_of_ecoregions, (-1, 1)))  # Transform the ecoregion labels using the OneHotEncoder
    X = scipy.sparse.hstack((X, list_of_ecoregions))  # Concatenate the ecoregion labels to the input data matrix
    return X, ohe_ecoregion, df_header

def add_dune(X, df_header, dunes, ohe_dune):
    """
    Add dune information to the input data matrix.

    Args:
        X (scipy.sparse._csr.csr_matrix): Input data matrix.
        df_header (pandas.core.frame.DataFrame): DataFrame containing header information.
        dunes (geopandas.GeoDataFrame): GeoDataFrame containing dune geometry information.
        ohe_dune (sklearn.preprocessing._encoders.OneHotEncoder): OneHotEncoder for dune labels.

    Returns:
        scipy.sparse._csr.csr_matrix: Updated data matrix.
        sklearn.preprocessing._encoders.OneHotEncoder: Updated OneHotEncoder for dune labels.
        pandas.core.frame.DataFrame: Updated DataFrame.
    """
    list_of_dunes = []  # Initialize a list to store dune labels
    for i in tqdm.tqdm(range(len(df_header)), desc="Dunes"):  # Iterate over each data point
        point = shapely.geometry.Point(df_header['Longitude'][i], df_header['Latitude'][i])  # Create a point object from longitude and latitude coordinates
        point = {'geometry': [point]}  # Create a GeoDataFrame with the point
        point = gpd.GeoDataFrame(point, crs="EPSG:4326")
        point = point.to_crs(crs=3857)   # Convert the point to the desired coordinate reference system
        point = shapely.geometry.Point(point.centroid.x, point.centroid.y)  # Calculate the centroid of the point
        list_of_dunes.append(point.within(dunes))  # Check if the point is within any dune
    list_of_dunes = ['Y_DUNES' if dune else 'N_DUNES' for dune in list_of_dunes]  # Map boolean values to dune labels ('Y_DUNES' or 'N_DUNES')
    if not ohe_dune:  # If OneHotEncoder for dune labels is not provided, initialize and fit a new one
        ohe_dune = sklearn.preprocessing.OneHotEncoder(handle_unknown='ignore')
        ohe_dune.fit(np.reshape(list_of_dunes, (-1, 1)))
    df_header['Dune'] = list_of_dunes  # Update the DataFrame with dune labels
    list_of_dunes = ohe_dune.transform(np.reshape(list_of_dunes, (-1, 1)))  # Transform dune labels using the OneHotEncoder
    X = scipy.sparse.hstack((X, list_of_dunes))  # Concatenate the dune labels to the input data matrix
    return X, ohe_dune, df_header

def add_coast(X, df_header, coasts, ohe_coast):
    """
    Add coast information to the input data matrix.

    Args:
        X (scipy.sparse._csr.csr_matrix): Input data matrix.
        df_header (pandas.core.frame.DataFrame): DataFrame containing header information.
        coasts (geopandas.GeoDataFrame): GeoDataFrame containing coast geometry information.
        ohe_coast (sklearn.preprocessing._encoders.OneHotEncoder): OneHotEncoder for coast labels.

    Returns:
        scipy.sparse._csr.csr_matrix: Updated data matrix.
        sklearn.preprocessing._encoders.OneHotEncoder: Updated OneHotEncoder for coast labels.
        pandas.core.frame.DataFrame: Updated DataFrame.
    """
    list_of_coasts = []  # Initialize a list to store coast labels
    for i in tqdm.tqdm(range(len(df_header)), desc="Coasts"):  # Iterate over each data point
        point = shapely.geometry.Point(df_header['Longitude'][i], df_header['Latitude'][i])  # Create a point object from longitude and latitude coordinates
        point = {'geometry': [point]}  # Create a GeoDataFrame with the point
        point = gpd.GeoDataFrame(point, crs="EPSG:4326")
        point = point.to_crs(crs=3857)  # Convert the point to the desired coordinate reference system
        point = shapely.geometry.Point(point.centroid.x, point.centroid.y)  # Calculate the centroid of the point
        if min(coasts.geometry.distance(point)) <= 5000:  # Calculate the distance between the point and each coast and check if the minimum distance to any coast is within a threshold (5000 units)
            list_of_coasts.append(coasts.geometry.distance(point).idxmin())  # Add the index of the coast with the minimum distance
        else:
            list_of_coasts.append("N_COAST")  # Add a label indicating no coast is within the threshold
    if not ohe_coast:  # If OneHotEncoder for coast labels is not provided, initialize and fit a new one
        ohe_coast = sklearn.preprocessing.OneHotEncoder(handle_unknown='ignore')
        ohe_coast.fit(np.reshape(list_of_coasts, (-1, 1)))
    df_header['Coast'] = list_of_coasts  # Update the DataFrame with coast labels
    list_of_coasts = ohe_coast.transform(np.reshape(list_of_coasts, (-1, 1)))  # Transform coast labels using the OneHotEncoder
    X = scipy.sparse.hstack((X, list_of_coasts))  # Concatenate the coast labels to the input data matrix
    return X, ohe_coast, df_header

data/preprocess_data.py
cesar-leblanc/hdm-framework
import pickle
import numpy as np
import scipy.sparse
import pandas as pd
import joblib
import torch

def set_le_header(args, le_header):
    """
    Save the label encoder for header information to a file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        le_header (sklearn.preprocessing._label.LabelEncoder): LabelEncoder for header information.
    """
    le_header_filepath = args.data_filepath + 'le_header.pkl'  # Define the file path to save the label encoder
    with open(le_header_filepath, 'wb') as f:  # Open the file in write binary mode and save the label encoder using pickle
        pickle.dump(le_header, f)

def set_eva_to_gbif_species(args, eva_to_gbif_species):
    """
    Save the mapping of EVA species to GBIF species to a file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        eva_to_gbif_species (dict): Mapping of EVA species to GBIF species.
    """
    eva_to_gbif_species_filepath = args.data_filepath + 'eva_to_gbif_species.pkl'  # Define the file path to save the mapping
    with open(eva_to_gbif_species_filepath, 'wb') as f:  # Open the file in write binary mode and save the mapping using pickle
        pickle.dump(eva_to_gbif_species, f)

def set_le_species(args, le_species):
    """
    Save the LabelEncoder object for species to a file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        le_species (sklearn.preprocessing._label.LabelEncoder): LabelEncoder object for species.
    """
    le_species_filepath = args.data_filepath + 'le_species.pkl'  # Define the file path to save the LabelEncoder object
    with open(le_species_filepath, 'wb') as f:  # Open the file in write binary mode and save the LabelEncoder object using pickle
        pickle.dump(le_species, f)

def set_split_assignment(args, split_assignment):
    """
    Save the split assignment array to a file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        split_assignment (numpy.ndarray): Array containing the split assignments.
    """
    split_assignment_filepath = args.data_filepath + 'split_assignments.npy'  # Define the file path to save the split assignment array
    np.save(split_assignment_filepath, split_assignment)  # Save the split assignment array using numpy's save function

def set_ohe_country(args, ohe_country):
    """
    Save the OneHotEncoder object for country to a file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        ohe_country (sklearn.preprocessing._encoders.OneHotEncoder): OneHotEncoder object for country.
    """
    ohe_country_filepath = args.data_filepath + 'ohe_country.pkl'  # Define the file path to save the OneHotEncoder object for country
    with open(ohe_country_filepath, 'wb') as f:  # Save the OneHotEncoder object using pickle
        pickle.dump(ohe_country, f)

def set_ohe_ecoregion(args, ohe_ecoregion):
    """
    Save the OneHotEncoder object for ecoregion to a file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        ohe_ecoregion (sklearn.preprocessing._encoders.OneHotEncoder): OneHotEncoder object for ecoregion.
    """
    ohe_ecoregion_filepath = args.data_filepath + 'ohe_ecoregion.pkl'  # Define the file path to save the OneHotEncoder object for ecoregion
    with open(ohe_ecoregion_filepath, 'wb') as f:  # Save the OneHotEncoder object using pickle
        pickle.dump(ohe_ecoregion, f)

def set_ohe_dune(args, ohe_dune):
    """
    Save the OneHotEncoder object for dune to a file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        ohe_dune (sklearn.preprocessing._encoders.OneHotEncoder): OneHotEncoder object for dune.
    """
    ohe_dune_filepath = args.data_filepath + 'ohe_dune.pkl'  # Define the file path to save the OneHotEncoder object for dune
    with open(ohe_dune_filepath, 'wb') as f:  # Save the OneHotEncoder object using pickles
        pickle.dump(ohe_dune, f)

def set_ohe_coast(args, ohe_coast):
    """
    Save the OneHotEncoder object for coast to a file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        ohe_coast (sklearn.preprocessing._encoders.OneHotEncoder): OneHotEncoder object for coast.
    """
    ohe_coast_filepath = args.data_filepath + 'ohe_coast.pkl'  # Define the file path to save the OneHotEncoder object for coast
    with open(ohe_coast_filepath, 'wb') as f:  # Save the OneHotEncoder object using pickle
        pickle.dump(ohe_coast, f)

def set_eva_header(args, eva_header):
    """
    Save the evaluation header DataFrame to a CSV file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        eva_header (pandas.core.frame.DataFrame): Evaluation header DataFrame.
    """
    eva_header_filepath = args.datasets_filepath + 'eva_header.csv'  # Define the file path to save the evaluation header CSV file
    eva_header.to_csv(eva_header_filepath, index=False)  # Save the evaluation header DataFrame to a CSV file

def set_eva_species(args, eva_species):
    """
    Save the evaluation species DataFrame to a CSV file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        eva_species (pandas.core.frame.DataFrame): Evaluation species DataFrame.
    """
    eva_species_filepath = args.datasets_filepath + 'eva_species.csv'  # Define the file path to save the evaluation species CSV file
    eva_species.to_csv(eva_species_filepath, index=False)  # Save the evaluation species DataFrame to a CSV file

def set_test_header(args, test_header):
    """
    Save the test header DataFrame to a CSV file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        test_header (pandas.core.frame.DataFrame): Test header DataFrame.
    """
    test_header_filepath = args.datasets_filepath + 'test_header.csv'  # Define the file path to save the test header CSV file
    test_header.to_csv(test_header_filepath, index=False)  # Save the test header DataFrame to a CSV file

def set_test_species(args, test_species):
    """
    Save the test species DataFrame to a CSV file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        test_species (pandas.core.frame.DataFrame): Test species DataFrame.
    """
    test_species_filepath = args.datasets_filepath + 'test_species.csv'  # Define the file path to save the test species CSV file
    test_species.to_csv(test_species_filepath, index=False)  # Save the test species DataFrame to a CSV file

def set_input_data(args, X):
    """
    Save the input data sparse matrix to a compressed NPZ file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        X (scipy.sparse._csr.csr_matrix): Input data sparse matrix.
    """
    input_data_filepath = args.data_filepath + 'input_data.npz'  # Define the file path to save the input data NPZ file
    scipy.sparse.save_npz(input_data_filepath, X)  # Save the input data sparse matrix to a compressed NPZ file
    
def set_test_data(args, X):
    """
    Save the test data sparse matrix to a compressed NPZ file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        X (scipy.sparse._csr.csr_matrix): Test data sparse matrix.
    """
    test_data_filepath = args.data_filepath + 'test_data.npz'  # Define the file path to save the test data NPZ file
    scipy.sparse.save_npz(test_data_filepath, X)  # Save the test data sparse matrix to a compressed NPZ file
    
def set_target_values(args, y):
    """
    Save the target values array to a NPY file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        y (np.ndarray): Target values array.
    """
    target_values_filepath = args.data_filepath + 'target_values.npy'  # Define the file path to save the target values NPY file
    np.save(target_values_filepath, y)  # Save the target values array to a NPY file
    
def set_scaler(args, scaler):
    """
    Save the scaler object to a pickle file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        scaler (sklearn.preprocessing._data.StandardScaler): Scaler object.
    """
    scaler_filepath = args.data_filepath + 'scaler.pkl'  # Define the file path to save the scaler pickle file
    with open(scaler_filepath, 'wb') as f:  # Save the scaler object to a pickle file
        pickle.dump(scaler, f)

def set_mlp_model(args, model):
    """
    Save the MLP model to a PyTorch model file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        model (models.mlp.MLP | torch.nn.parallel.data_parallel.DataParallel): MLP model.
    """
    mlp_filepath = args.models_filepath + 'MLP.pth'  # Define the file path to save the MLP model file
    torch.save(model, mlp_filepath)  # Save the MLP model to a PyTorch model file
    
def set_rfc_model(args, model):
    """
    Save the RFC model to a joblib file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        model (models.rfc.RFC): RFC model.
    """
    rfc_filepath = args.models_filepath + 'RFC.joblib'  # Define the file path to save the RFC model file
    joblib.dump(model, rfc_filepath, compress=True)  # Save the RFC model to a joblib file
    
def set_xgb_model(args, model):
    """
    Save the XGB model to a JSON file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        model (models.xgb.XGB): XGB model.
    """
    xgb_filepath = args.models_filepath + 'XGB.json'  # Define the file path to save the XGB model file
    model.save_model(xgb_filepath)  # Save the XGB model to a JSON file

def set_tnc_model(args, model):
    """
    Save the TNC model to a file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        model (models.tnc.TNC): TNC model.
    """
    tnc_filepath = args.models_filepath + 'TNC'  # Define the file path to save the TNC model file
    model.save_model(tnc_filepath)  # Save the TNC model to a file
    
def set_ftt_model(args, model):
    """
    Save the FTT model to a PyTorch model file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        model (models.ftt.FTT | torch.nn.parallel.data_parallel.DataParallel): FTT model.
    """
    ftt_filepath = args.models_filepath + 'FTT.pth'  # Define the file path to save the FTT model file
    torch.save(model, ftt_filepath)  # Save the FTT model to a PyTorch model file

def set_predictions(args, predictions):
    """
    Save the predictions array to a text file.

    Args:
        args (argparse.Namespace): Arguments passed to the function.
        predictions (np.ndarray): Predictions array.
    """
    predictions_filepath = args.data_filepath + 'predictions.txt'  # Define the file path to save the predictions text file
    np.savetxt(predictions_filepath, predictions, fmt="%s")  # Save the predictions array to a text file

data/save_data.py
cesar-leblanc/hdm-framework
import argparse

from cli import add_all_parsers
from utils import disable_warnings, print_parameters, set_seed, disable_bytecode_generation

from pipelines.check import Check
from pipelines.evaluation import Evaluation
from pipelines.dataset import Dataset
from pipelines.training import Training
from pipelines.prediction import Prediction

AVAILABLE_PIPELINES = {  # Dictionary of all available pipelines
    "check": Check,
    "dataset": Dataset,
    "evaluation": Evaluation,
    "training": Training,
    "prediction": Prediction
}

def run(args):
    """
    Run the framework for Habitat Distribution Modeling.

    Args:
        args (argparse.Namespace): Parsed command-line arguments.
    """
    disable_warnings(args)  # Ignore all warnings during the code's execution
    disable_bytecode_generation(args)  # Disable bytecode generation to improve startup performance
    print_parameters(args)  # Print the framework parameters
    set_seed(args)  # Set the random seed for reproducibility
    pipeline = AVAILABLE_PIPELINES[args.pipeline]()  # Create an instance of the selected pipeline
    pipeline.run(args)  # Run the pipeline with the provided arguments

if __name__ == "__main__":
    parser = argparse.ArgumentParser()  # Create an argument parser
    add_all_parsers(parser)  # Add parsers for all the available options
    args = parser.parse_args()  # Parse the command-line arguments
    run(args)  # Run the framework with the parsed arguments

main.py
cesar-leblanc/hdm-framework
import rtdl
import torch
import numpy as np
import time
import scipy.special

from utils import get_device, get_criterion, get_optimizer, get_scheduler, get_metric, NormedLinear, get_model_parameters

from data.load_data import get_split_assignments, get_eunis_red_list_crosswalks, get_endangered_red_list_habitats, get_endangered_eunis_habitats, get_ftt_model, get_scaler
from data.preprocess_data import add_fold_assignments, add_dataloader, add_standardization, add_normalization, add_dropout, add_binarization, add_log, add_noise, add_rank, remove_features, add_augmentation, add_endangered_habitats, add_predictions_decoding
from data.save_data import set_scaler, set_ftt_model, set_predictions

class FTT(rtdl.FTTransformer):
    """
    Feature Transformer for Habitat Distribution Modeling.
    """
    def __init__(self, n_num_features=None, cat_cardinalities=None, d_token=16, n_blocks=1, attention_n_heads=4, attention_dropout=0.3, attention_initialization='kaiming', attention_normalization='LayerNorm', ffn_d_hidden=16, ffn_dropout=0.1, ffn_activation='ReGLU', ffn_normalization='LayerNorm', residual_dropout=0.0, prenormalization=True, first_prenormalization=False, last_layer_query_idx=[-1], n_tokens=None, kv_compression_ratio=0.004, kv_compression_sharing='headwise', head_activation='ReLU', head_normalization='LayerNorm', d_out=None):
        """
        Initialize the Feature Transformer.

        Args:
            n_num_features (int, optional): The number of input features. Defaults to None.
            cat_cardinalities (list, optional): The number of unique values for each feature. Defaults to None.
            d_token (int, optional): The size of one token. Defaults to 16.
            n_blocks (int, optional): The number of Transformer blocks.  Defaults to 1.
            attention_n_heads (int, optional): The number of attention heads. Defaults to 4.
            attention_dropout (float, optional): The dropout for attention blocks. Defaults to 0.3.
            attention_initialization (str, optional): The initialization for attention blocks. Defaults to 'kaiming'.
            attention_normalization (str, optional): The normalization for attention blocks. Defaults to 'LayerNorm'.
            ffn_d_hidden (int, optional): The input size for the second linear layer in the Feed-Forward Network module. Defaults to 16.
            ffn_dropout (float, optional): The dropout rate after the first linear layer in the Feed-Forward Network module. Defaults to 0.1.
            ffn_activation (str, optional): The activation used in the Feed-Forward Network. Defaults to 'ReGLU'.
            ffn_normalization (str, optional): The normalization used in the Feed-Forward Network. Defaults to 'LayerNorm'.
            residual_dropout (float, optional): The dropout rate for the output of each residual branch of all Transformer blocks. Defaults to 0.0.
            prenormalization (bool, optional): The choice to place normalizations at the beginning of each residual branch. Defaults to True.
            first_prenormalization (bool, optional): The choice to keep the first normalization from the first Transformer layer. Defaults to False.
            last_layer_query_idx (list, optional): The indices of tokens that should be processed by the last Transformer block. Defaults to [-1].
            n_tokens (int, optional): The option for fast linear attention. Defaults to None.
            kv_compression_ratio (float, optional): The choice to apply a technique to speed up attention modules when the number of features is large. Defaults to 0.004.
            kv_compression_sharing (str, optional): Weight sharing policy for the technique to speed up attention modules when the number of features is large. Defaults to 'headwise'.
            head_activation (str, optional): The activation used in the heads. Defaults to 'ReLU'.
            head_normalization (str, optional): The normalization used in the heads. Defaults to 'LayerNorm'.
            d_out (int, optional): The number of output classes. Defaults to None.
        """
        feature_tokenizer = rtdl.FeatureTokenizer(  # The odule combining `NumericalFeatureTokenizer` and `CategoricalFeatureTokenizer`
            n_num_features=n_num_features,  # The number of input features
            cat_cardinalities=cat_cardinalities,  # The number of unique values for each feature
            d_token=d_token  # The size of one token
        )
        transformer = rtdl.Transformer(  # The Transformer with extra features
            d_token=d_token,  # The size of one token
            n_blocks=n_blocks,  # The number of Transformer blocks
            attention_n_heads=attention_n_heads,  # The number of attention heads
            attention_dropout=attention_dropout,  # The dropout for attention blocks
            attention_initialization=attention_initialization,  # The initialization for attention blocks
            attention_normalization=attention_normalization,  # The normalization for attention blocks
            ffn_d_hidden=ffn_d_hidden,  # The input size for the second linear layer in the Feed-Forward Network module
            ffn_dropout=ffn_dropout,  # The dropout rate after the first linear layer in the Feed-Forward Network module
            ffn_activation=ffn_activation,  # The activation used in the Feed-Forward Network
            ffn_normalization=ffn_normalization,  # The normalization used in the Feed-Forward Network
            residual_dropout=residual_dropout,  # The dropout rate rate for the output of each residual branch of all Transformer blocks
            prenormalization=prenormalization,  # The choice to place normalizations at the beginning of each residual branch
            first_prenormalization=first_prenormalization,  # The choice to keep the first normalization from the first Transformer layer
            last_layer_query_idx=last_layer_query_idx,  # The indices of tokens that should be processed by the last Transformer block
            n_tokens=n_tokens,  # The option for fast linear attention
            kv_compression_ratio=kv_compression_ratio,  # The choice to apply a technique to speed up attention modules when the number of features is large
            kv_compression_sharing=kv_compression_sharing,  # Weight sharing policy for the technique to speed up attention modules when the number of features is large
            head_activation=head_activation,  # The activation used in the heads
            head_normalization=head_normalization,  # The normalization used in the heads
            d_out=d_out  # The number of output classes
        )
        super(FTT, self).__init__(feature_tokenizer, transformer)  # Initialize the parent class, `rtdl.FTTransformer`
        
    def evaluate_model(self, args, X, y, le_species, le_header):
        """
        Evaluate the model using the specified arguments and input data.

        Args:
            args (argparse.Namespace): Arguments for model evaluation.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing.LabelEncoder): LabelEncoder for species.
            le_header (sklearn.preprocessing.LabelEncoder): LabelEncoder for header.

        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank-based features to the input data
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features to the input data
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features to the input data
        std_acc = []
        avg_acc = 0.0  # Initialize variables to track the average accuracy
        n_samples = 0  # Initialize variables to track the number of samples
        split_assignments = get_split_assignments(args)
        best_accuracy = -np.inf
        best_fold = -1
        start_evaluation = time.time()
        for fold in range(args.n_folds):  # Perform evaluation for each fold
            best_fold_accuracy = self.evaluate_fold(args, X, y, le_species, le_header, fold, split_assignments)
            if best_fold_accuracy > best_accuracy:
                best_accuracy = best_fold_accuracy
                best_fold = fold
                print("\nFold {} did improve the model.".format(best_fold))
            else:
                print("\nFold {} did not improve the model.".format(fold))
            std_acc.append(best_fold_accuracy)
            avg_acc += best_fold_accuracy  # Accumulate the accuracy
            n_samples += 1  # Accumulate the number of samples
        end_evaluation = time.time()
        time_evaluation = end_evaluation - start_evaluation
        avg_acc /= n_samples  # Compute the average accuracy
        std_acc = torch.std(torch.FloatTensor(std_acc))
        print(f"\nMean {args.metric} and standard deviation after {args.n_folds} folds of {args.num_epochs} epochs: {avg_acc:.4f}% & {std_acc:.4f}.")  # Print evaluation results
        print(f"Total time of the evaluation: {time_evaluation:.2f}s.")  # Print evaluation time
        
    def train_model(self, args, X, y, le_species, le_header):
        """
        Train the model using the specified arguments and input data.

        Args:
            args (argparse.Namespace): Arguments for model training.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing.LabelEncoder): LabelEncoder for species.
            le_header (sklearn.preprocessing.LabelEncoder): LabelEncoder for header.

        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank-based features to the input data
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features to the input data
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features to the input data
        if args.endangered:
            eunis_red_list_crosswalks = get_eunis_red_list_crosswalks(args)  # Get endangered habitats
            endangered_red_list_habitats = get_endangered_red_list_habitats(args)
            endangered_eunis_habitats = get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header)
            X, y = add_endangered_habitats(X, y, endangered_eunis_habitats)  # Modify input data and labels accordingly
        if args.augmentation:
            if not args.endangered:
                labels = None
            else:
                labels = endangered_eunis_habitats
            X, y = add_augmentation(X, y, labels)  # Add data augmentation to the input data and labels
        if args.dropout > 0:
            X = add_dropout(X, args.dropout, len(le_species.classes_))  # Add dropout to the input data
        if args.noise:
            X = add_noise(X)  # Add noise to the input data
        if args.log:
            X = add_log(X, len(le_species.classes_))  # Apply logarithmic transformation to the input data
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X = remove_features(args, X, features)  # Remove specific features from the input data
        if args.standardization:
            X, scaler = add_standardization(X, scaler=None)  # Add standardization to the input data and save the scaler
            set_scaler(args, scaler)
        train_dataloader = add_dataloader(args, X, y, shuffle=True)
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))
        model = FTT(**model_parameters)
        if args.criterion == 'imbal-noised-top-k':
            LastLinearLayer = NormedLinear
            model.transformer.head.linear = LastLinearLayer(model.transformer.head.linear.in_features, model.transformer.head.linear.out_features)
        device = get_device(args)  # Set the device to run the model on to be GPU
        if args.data_parallelism:
            model = torch.nn.DataParallel(model)  # Run the model parallelly
        model.to(device)  # Use GPU for model
        criterion = get_criterion(args, y)  # Instantiate loss class
        criterion.to(device)
        optimizer = get_optimizer(args, model)  # Instantiate optimizer class
        scheduler = get_scheduler(args, optimizer)  # Instantiate step learning scheduler class
        best_accuracy = -np.inf
        best_epoch = -1
        best_model_state = model.state_dict()  # Set up a variable to store the best model's state
        no_improvement_epochs = 0  # Set up a counter for the number of epochs without improvement
        start_training = time.time()
        for epoch in range(args.num_epochs):
            epoch_accuracy, model = self.train_epoch(args, model, train_dataloader, epoch, device, optimizer, criterion, le_header)
            scheduler.step(epoch_accuracy)  # Decay Learning Rate, pass validation accuracy for tracking at every epoch
            if epoch_accuracy > best_accuracy:
                best_accuracy = epoch_accuracy
                best_epoch = epoch
                best_model_state = model.state_dict()
                set_ftt_model(args, best_model_state)
                print("\nSuccessfully saved model at Models/FTT.pth")
                no_improvement_epochs = 0
            else:  # Check if the accuracy is not the best seen so far
                no_improvement_epochs += 1  # Increment the number of epochs without improvement
                print("\nEpoch {} did not improve the model.".format(epoch))
            if no_improvement_epochs == args.num_iter_no_change:  # Check if the model's accuracy on the test set has not improved for the num_iter_no_change number of epochs               
                print(f'\nEarly stopping occurred at epoch {epoch} with best epoch = {best_epoch} and best {args.metric} = {best_accuracy:.4f}%.')
                break  # Stop the training process
        end_training = time.time()
        time_training = end_training - start_training
        if no_improvement_epochs != args.num_iter_no_change:
            print(f'\nEarly stopping dit not occur as best epoch = {best_epoch} and best {args.metric} = {best_accuracy:.4f}%.')
        print(f"Total time of the training: {time_training:.2f}s.")

    def predict_model(self, args, X, y, le_species, le_header):
        """
        Predict the habitat labels for given input data using a trained model.

        Args:
            args (argparse.Namespace): Command-line arguments.
            X (scipy.sparse._csr.csr_matrix): Input data features.
            y (numpy.ndarray): Input data labels.
            le_species (sklearn.preprocessing.LabelEncoder): LabelEncoder for species classes.
            le_header (sklearn.preprocessing.LabelEncoder): LabelEncoder for habitat classes.

        Returns:
            None
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank features to input data
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features to input data
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features to input data
        if args.log:
            X = add_log(X, len(le_species.classes_))  # Add logarithmic features to input data
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X = remove_features(args, X, features)  # Remove specified features from input data
        if args.standardization:
            scaler = get_scaler(args)
            X, _ = add_standardization(X, scaler=scaler)  # Add standardized features to input data using the scaler
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Get model parameters
        model = FTT(**model_parameters)  # Create an instance of the FTT model
        if args.criterion == 'imbal-noised-top-k':
            LastLinearLayer = NormedLinear
            model.transformer.head.linear = LastLinearLayer(model.transformer.head.linear.in_features, model.transformer.head.linear.out_features)  # Replace the last linear layer with a custom NormedLinear layer
        device = get_device(args)  # Set the device to run the model on to be GPU
        if args.data_parallelism:
            model = torch.nn.DataParallel(model)  # Run the model parallelly
        model = get_ftt_model(args, model)  # Load the trained model parameters
        model.to(device)  # Move the model to the specified device (GPU or CPU)
        X = torch.from_numpy(X.toarray()).float()  # Convert input data to a PyTorch tensor
        model.eval()  # Set the model to evaluation mode
        start_prediction = time.time()  # Start timing the prediction process
        with torch.no_grad():  # Turn off gradient computation to save memory
            predictions = model(x_num=X, x_cat=None)  # Pass the new data through the model
        end_prediction = time.time()  # End timing the prediction process
        time_prediction = end_prediction - start_prediction  # Compute the time taken for prediction
        predictions = scipy.special.softmax(predictions.cpu(), axis=1)  # Apply softmax function to convert logits to probabilities
        predictions = [np.argmax(sub_arr) for sub_arr in predictions]  # Get the indices of the maximum probability for each prediction
        predictions = add_predictions_decoding(predictions, le_header)  # Decode the predicted labels using the LabelEncoder for habitats
        set_predictions(args, predictions)  # Store the predictions in the specified location
        print("\nSuccessfully saved predictions at Data/predictions.txt")
        print(f"Total time of the prediction: {time_prediction:.2f}s.")  # Print the total time taken for prediction
    
    def evaluate_fold(self, args, X, y, le_species, le_header, fold, split_assignments):
        """
        Evaluate a single fold of the dataset for Habitat Distribution Modeling.

        Args:
            args (argparse.Namespace): Command-line arguments.
            X (scipy.sparse._csr.csr_matrix): Input data features.
            y (numpy.ndarray): Input data labels.
            le_species (sklearn.preprocessing.LabelEncoder): LabelEncoder for species classes.
            le_header (sklearn.preprocessing.LabelEncoder): LabelEncoder for habitat classes.
            fold (int): Fold number for evaluation.
            split_assignments (numpy.ndarray): Array indicating the fold assignments for each sample.

        Returns:
            float: Best accuracy achieved in the fold.
        """
        print('\n' + '*'*11)
        print('* Fold: {} *'.format(fold))  # Print fold
        print('*'*11 + '\n' + '-'*30)
        X_train, X_test, y_train, y_test = add_fold_assignments(X, y, split_assignments, fold)  # Split data into training and test sets based on fold assignments
        if args.endangered:
            eunis_red_list_crosswalks = get_eunis_red_list_crosswalks(args)
            endangered_red_list_habitats = get_endangered_red_list_habitats(args)
            endangered_eunis_habitats = get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header)  # Get endangered habitats
            X_test, y_test = add_endangered_habitats(X_test, y_test, endangered_eunis_habitats)  # Add endangered habitats to the test set
        if args.augmentation:
            if not args.endangered:
                labels = None
            else:
                labels = endangered_eunis_habitats
            X_train, y_train = add_augmentation(X_train, y_train, labels)  # Augment the training set with additional samples
        if args.dropout > 0:
            X_test = add_dropout(X_test, args.dropout, len(le_species.classes_))  # Apply dropout to the test set
        if args.noise:
            X_test = add_noise(X_test)  # Add noise to the test set
        if args.log:
            X_train = add_log(X_train, len(le_species.classes_))  # Apply logarithmic transformation to the training set
            X_test = add_log(X_test, len(le_species.classes_))  # Apply logarithmic transformation to the test set
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X_train = remove_features(args, X_train, features)  # Remove specific features from the training set
            X_test = remove_features(args, X_test, features)  # Remove specific features from the test set
        if args.standardization:
            X_train, scaler = add_standardization(X_train, scaler=None)  # Apply standardization to the training set
            X_test, scaler = add_standardization(X_test, scaler=scaler)  # Apply standardization to the test set
        train_dataloader = add_dataloader(args, X_train, y_train, shuffle=True)  # Create a dataloader for the training set
        test_dataloader = add_dataloader(args, X_test, y_test, shuffle=False)  # Create a dataloader for the test set
        model_parameters = get_model_parameters(args, X_train.shape[1], len(le_header.classes_))   # Get model parameters
        model = FTT(**model_parameters)  # Create an instance of the FTT model
        if args.criterion == 'imbal-noised-top-k':
            LastLinearLayer = NormedLinear
            model.transformer.head.linear = LastLinearLayer(model.transformer.head.linear.in_features, model.transformer.head.linear.out_features)  # Replace the last linear layer with a customized linear layer
        device = get_device(args)  # Set the device to run the model on to be GPU
        if args.data_parallelism:
            model = torch.nn.DataParallel(model)  # Run the model parallelly
        model.to(device)  # Move the model to the specified device (GPU or CPU)
        criterion = get_criterion(args, y_train)  # Instantiate loss class
        criterion.to(device)  # Move the loss function to the specified device (GPU or CPU)
        optimizer = get_optimizer(args, model)  # Instantiate optimizer class
        scheduler = get_scheduler(args, optimizer)  # Instantiate step learning scheduler class
        best_fold_accuracy = -np.inf  # Set up a variable to store the best accuracy
        best_epoch = -1  # Set up a variable to store the best epoch
        no_improvement_epochs = 0  # Set up a counter for the number of epochs without improvement
        start_fold = time.time()  # Start timing the evaluation of the fold
        for epoch in range(args.num_epochs):
            accuracy, model = self.evaluate_epoch(args, model, train_dataloader, test_dataloader, epoch, device, optimizer, criterion, le_header)
            scheduler.step(accuracy)  # Decay Learning Rate, pass validation accuracy for tracking at every epoch
            if accuracy > best_fold_accuracy:  # Check if the model's state has the highest accuracy on the test set so far
                best_fold_accuracy = accuracy  # Update the best accuracy
                best_epoch = epoch  # Update the best epoch
                no_improvement_epochs = 0  # Reset the number of epochs without improvement
            else:  # Check if the test accuracy is not the best seen so far
                no_improvement_epochs += 1  # Increment the number of epochs without improvement
            if no_improvement_epochs == args.num_iter_no_change:  # Check if the model's accuracy on the test set has not improved for the num_iter_no_change number of epochs               
                print(f'\nEarly stopping occurred at epoch {epoch} with best epoch = {best_epoch} and best {args.metric} = {best_fold_accuracy}')
                break  # Stop the training process
        end_fold = time.time()  # End timing the evaluation of the fold
        time_fold = end_fold - start_fold  # Calculate the time taken for the evaluation of the fold
        if no_improvement_epochs != args.num_iter_no_change:
            print(f'\nEarly stopping dit not occur as best epoch = {best_epoch} and best {args.metric} = {best_fold_accuracy}')
        print(f"Total time for the evaluation of the fold: {time_fold:.2f}s.")
        return best_fold_accuracy
    
    def train_epoch(self, args, model, train_dataloader, epoch, device, optimizer, criterion, le_header):
        """
        Train the model for a single epoch.

        Args:
            args (argparse.Namespace): Arguments passed to the training process.
            model (models.ftt.FTT | torch.nn.parallel.data_parallel.DataParallel): The deep learning model.
            train_dataloader (torch.utils.data.dataloader.DataLoader): Dataloader for the training dataset.
            epoch (int): Current epoch number.
            device (torch.device): Device to run the model on (GPU or CPU).
            optimizer (torch.optim.sgd.SGD | torch.optim.adam.Adam | torch.optim.adamw.AdamW): Optimizer for updating model parameters.
            criterion (torch.nn.modules.loss.CrossEntropyLoss | pytopk.noised_losses.BalNoisedTopK | pytopk.noised_losses.ImbalNoisedTopK): Loss function.
            le_header (sklearn.preprocessing.LabelEncoder): Label encoder for habitat classes.

        Returns:
            tuple: A tuple containing the accuracy of the epoch and the updated model.
        """
        start_epoch = time.time()  # Start timing the epoch
        for batch, labels in train_dataloader:  # Iterate through train dataset
            batch = batch.requires_grad_().to(device)  # Load batches with gradient accumulation capabilities
            labels = labels.to(device)  # Use GPU for tensors
            optimizer.zero_grad()  # Clear gradients w.r.t. parameters
            outputs = model(x_num=batch, x_cat=None)  # Forward pass to get output/logits
            loss = criterion(outputs, labels)  # Calculate Loss: softmax --> cross entropy loss
            loss.backward()  # Getting gradients w.r.t. parameters
            optimizer.step()  # Updating parameters
        end_epoch = time.time()  # End timing the epoch
        time_epoch = end_epoch - start_epoch  # Calculate the time taken for the epoch
        metric = get_metric(args, len(le_header.classes_))  # Instantiate the accuracy metric
        metric.to(device)  # Move the metric to the device
        model.eval()  # Set the model to evaluation mode
        with torch.no_grad():  # Disable gradient computation
            for batch, labels in train_dataloader:  # Iterate through the batches in the training dataset
                batch = batch.to(device)  # Move the batch to the device
                labels = labels.to(device)  # Move the labels to the device
                outputs = model(x_num=batch, x_cat=None)  # Forward pass to get output/logits
                accuracy = metric(outputs, labels)  # Calculate the accuracy
        accuracy_epoch = metric.compute() * 100  # Compute the overall accuracy for the epoch
        print('\n' + '-'*30)
        print('Epoch {} completed.'.format(epoch))  # Print Epoch
        print('LR: {}.'.format(optimizer.param_groups[0]['lr']))  # Print Learning Rate
        print(f"Time: {time_epoch:.2f}s.")  # Print Time
        print(f'Accuracy: {accuracy_epoch   :.4f}%.')  # Print Accuracy
        print('-'*30)
        model.train()  # Set the model back to training mode
        return accuracy_epoch, model
    
    def evaluate_epoch(self, args, model, train_dataloader, test_dataloader, epoch, device, optimizer, criterion, le_header):
        """
        Evaluate a single epoch of the training process.

        Args:
            args (argparse.Namespace): Training arguments.
            model (models.ftt.FTT | torch.nn.parallel.data_parallel.DataParallel): The deep learning model.
            train_dataloader (torch.utils.data.dataloader.DataLoader): Dataloader for the training dataset.
            test_dataloader (torch.utils.data.dataloader.DataLoader): Dataloader for the test dataset.
            epoch (int): Current epoch number.
            device (torch.device): Device to perform computations on (e.g., CPU or GPU).
            optimizer (torch.optim.sgd.SGD | torch.optim.adam.Adam | torch.optim.adamw.AdamW): Optimizer for updating model parameters.
            criterion (torch.nn.modules.loss.CrossEntropyLoss | pytopk.noised_losses.BalNoisedTopK | pytopk.noised_losses.ImbalNoisedTopK): Loss function criterion.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for the header.

        Returns:
            tuple: A tuple containing the accuracy of the epoch and the updated model.
        """
        start_epoch = time.time()  # Start timing the epoch
        for batch, labels in train_dataloader:  # Iterate through train dataset
            batch = batch.requires_grad_().to(device)  # Load batches with gradient accumulation capabilities
            labels = labels.to(device)  # Use GPU for tensors
            optimizer.zero_grad()  # Clear gradients w.r.t. parameters
            outputs = model(x_num=batch, x_cat=None)  # Forward pass to get output/logits
            loss = criterion(outputs, labels)  # Calculate Loss: softmax --> cross entropy loss
            loss.backward()  # Getting gradients w.r.t. parameters
            optimizer.step()  # Updating parameters
        end_epoch = time.time()  # End timing the epoch
        time_epoch = end_epoch - start_epoch  # Calculate the time taken for the epoch
        metric = get_metric(args, len(le_header.classes_))  # Instantiate the accuracy metric
        metric.to(device)  # Move the metric to the device
        model.eval()  # Set the model to evaluation mode
        with torch.no_grad():  # Disable gradient computation
            for batch, labels in test_dataloader:  # Iterate through the batches in the test dataset
                batch = batch.to(device)  # Move the batch to the device
                labels = labels.to(device)  # Move the labels to the device
                outputs = model(x_num=batch, x_cat=None)  # Forward pass to get output/logits
                accuracy = metric(outputs, labels)  # Calculate the accuracy
        accuracy_epoch = metric.compute() * 100  # Compute the overall accuracy for the epoch
        print('Epoch {} completed.'.format(epoch))  # Print Epoch
        print('LR: {}.'.format(optimizer.param_groups[0]['lr']))  # Print Learning Rate
        print(f"Time: {time_epoch:.2f}s.")  # Print Time
        print(f'Accuracy: {accuracy_epoch:.4f}%.')  # Print Accuracy
        print('-'*30)
        model.train()  # Set the model back to training mode
        return accuracy_epoch, model

    def run(self, args, X, y, le_species, le_header):
        """
        Run the FTT model based on the specified pipeline.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.pipeline == 'evaluation':
            self.evaluate_model(args, X, y, le_species, le_header)  # Run evaluation pipeline
        elif args.pipeline == 'training':
            self.train_model(args, X, y, le_species, le_header)  # Run training pipeline
        elif args.pipeline == 'prediction':
            self.predict_model(args, X, y, le_species, le_header)  # Run prediction pipeline

models/ftt.py
cesar-leblanc/hdm-framework
import torch
import numpy as np
import time
import scipy.special

from utils import get_device, get_criterion, get_optimizer, get_scheduler, get_metric, NormedLinear, get_model_parameters

from data.load_data import get_split_assignments, get_eunis_red_list_crosswalks, get_endangered_red_list_habitats, get_endangered_eunis_habitats, get_scaler, get_mlp_model
from data.preprocess_data import add_fold_assignments, add_dataloader, add_standardization, add_normalization, add_dropout, add_binarization, add_log, add_noise, add_rank, remove_features, add_augmentation, add_endangered_habitats, add_predictions_decoding
from data.save_data import set_scaler, set_mlp_model, set_predictions

class MLP(torch.nn.Module):
    """
    Multi-Layer Perceptron (MLP) model for habitat distribution modeling.
    """
    def __init__(self, input_dim=None, num_hidden_layers=1, num_neurons=200, output_dim=None, activation_func='relu', weight_init=None, dropout_p=0, use_batch_norm=False):
        """
        Initialize the MLP model.

        Args:
            input_dim (int, optional): Dimensionality of the input features. Defaults to None.
            num_hidden_layers (int, optional): Number of hidden layers. Defaults to 1.
            num_neurons (int, optional): Number of neurons in the hidden layers. Defaults to 200.
            output_dim (int, optional): Dimensionality of the output. Defaults to None.
            activation_func (str, optional): Activation function to use. Defaults to 'relu'.
            weight_init (str, optional): Weight initialization method.  Defaults to None.
            dropout_p (float, optional): Dropout probability. Defaults to 0.
            use_batch_norm (bool, optional): Whether to use batch normalization or not. Defaults to False.
        """
        super(MLP, self).__init__()  # Initialize the parent class, `torch.nn.Module`
        self.fc1 = torch.nn.Linear(input_dim, num_neurons)  # Create a linear layer with `input_dim` input units and `num_neurons` output units
        self.use_batch_norm = use_batch_norm  # Save whether batch normalization should be used or not
        if self.use_batch_norm:
            self.bn1 = torch.nn.BatchNorm1d(num_neurons)  # Create a batch normalization layer with `num_neurons` input units
        if activation_func == 'relu':  # Determine which activation function to use based on the `activation_func` parameter
            self.activation1 = torch.nn.ReLU()
        elif activation_func == 'sigmoid':
            self.activation1 = torch.nn.Sigmoid()
        elif activation_func == 'tanh':
            self.activation1 = torch.nn.Tanh()
        self.dropout1 = torch.nn.Dropout(p=dropout_p)  # Create a dropout layer with a probability of dropping out `dropout_p`

        self.hidden_layers = torch.nn.ModuleList()  # Create an empty list of hidden layers
        for i in range(num_hidden_layers - 1):  # Iterate through the number of hidden layers
            self.hidden_layers.append(torch.nn.Linear(num_neurons, num_neurons))  # Add a linear layer with `num_neurons` input and output units to the list
            if self.use_batch_norm:
                self.hidden_layers.append(torch.nn.BatchNorm1d(num_neurons))  # Add a batch normalization layer with `num_neurons` input units to the list
            if activation_func == 'relu':  # Determine which activation function to use based on the `activation_func` parameter
                self.hidden_layers.append(torch.nn.ReLU())  # Add a ReLU activation function to the list
            elif activation_func == 'sigmoid':
                self.hidden_layers.append(torch.nn.Sigmoid())  # Add a sigmoid activation function to the list
            elif activation_func == 'tanh':
                self.hidden_layers.append(torch.nn.Tanh())  # Add a tanh activation function to the list
            self.hidden_layers.append(torch.nn.Dropout(p=dropout_p))  # Add a dropout layer with a probability of dropping out `dropout_p` to the list

        self.fc2 = torch.nn.Linear(num_neurons, output_dim)  # Create a linear layer with `num_neurons` input units and `output_dim` output units

        if weight_init == 'xavier':  # If Xavier initialization is selected
            torch.nn.init.xavier_normal_(self.fc1.weight)  # Initialize the weights of the first linear layer using Xavier initialization
            torch.nn.init.xavier_normal_(self.fc2.weight)  # Initialize the weights of the second linear layer using Xavier initialization
            for layer in self.hidden_layers:  # Iterate through the hidden layers
                if isinstance(layer, torch.nn.Linear):  # If the layer is a linear layer
                    torch.nn.init.xavier_normal_(layer.weight)  # Initialize the weights using Xavier initialization
        elif weight_init == 'kaiming':  # If Kaiming initialization is selected
            torch.nn.init.kaiming_normal_(self.fc1.weight, nonlinearity='relu')  # Initialize the weights of the first linear layer using Kaiming initialization with a ReLU nonlinearity
            torch.nn.init.kaiming_normal_(self.fc2.weight, nonlinearity='relu')  # Initialize the weights of the second linear layer using Kaiming initialization with a ReLU nonlinearity
            for layer in self.hidden_layers:  # Iterate through the hidden layers
                if isinstance(layer, torch.nn.Linear):  # If the layer is a linear layer
                    torch.nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')  # Initialize the weights using Kaiming initialization

    def forward(self, x):
        """
        Forward pass of the HabitatModel.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor.
        """
        x = self.fc1(x)
        if self.use_batch_norm:
            x = self.bn1(x)
        x = self.activation1(x)
        x = self.dropout1(x)
        for layer in self.hidden_layers:
            x = layer(x)
        x = self.fc2(x)
        return x

    def evaluate_model(self, args, X, y, le_species, le_header):
        """
        Evaluate the model using the specified arguments and input data.

        Args:
            args (argparse.Namespace): Arguments for model evaluation.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): LabelEncoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): LabelEncoder for header.
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank-based features to the input data
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features to the input data
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features to the input data
        std_acc = []
        avg_acc = 0.0  # Initialize variables to track the average accuracy
        n_samples = 0  # Initialize variables to track the number of samples
        split_assignments = get_split_assignments(args)
        best_accuracy = -np.inf
        best_fold = -1
        start_evaluation = time.time()
        for fold in range(args.n_folds):  # Perform evaluation for each fold
            best_fold_accuracy = self.evaluate_fold(args, X, y, le_species, le_header, fold, split_assignments)
            if best_fold_accuracy > best_accuracy:
                best_accuracy = best_fold_accuracy
                best_fold = fold
                print("\nFold {} did improve the model.".format(best_fold))
            else:
                print("\nFold {} did not improve the model.".format(fold))
            std_acc.append(best_fold_accuracy)
            avg_acc += best_fold_accuracy  # Accumulate the accuracy
            n_samples += 1  # Accumulate the number of samples
        end_evaluation = time.time()
        time_evaluation = end_evaluation - start_evaluation
        avg_acc /= n_samples  # Compute the average accuracy
        std_acc = torch.std(torch.FloatTensor(std_acc))
        print(f"\nMean {args.metric} and standard deviation after {args.n_folds} folds of {args.num_epochs} epochs: {avg_acc:.4f}% & {std_acc:.4f}.")  # Print evaluation results
        print(f"Total time of the evaluation: {time_evaluation:.2f}s.")  # Print evaluation time

    def train_model(self, args, X, y, le_species, le_header):
        """
        Train the model using the specified arguments and input data.

        Args:
            args (argparse.Namespace): Arguments for model training.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): LabelEncoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): LabelEncoder for header.
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank-based features to the input data
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features to the input data
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features to the input data
        if args.endangered:
            eunis_red_list_crosswalks = get_eunis_red_list_crosswalks(args)
            endangered_red_list_habitats = get_endangered_red_list_habitats(args)  # Get endangered habitats
            endangered_eunis_habitats = get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header)
            X, y = add_endangered_habitats(X, y, endangered_eunis_habitats)  # Modify input data and labels accordingly
        if args.augmentation:
            if not args.endangered:
                labels = None
            else:
                labels = endangered_eunis_habitats  
            X, y = add_augmentation(X, y, labels)  # Add data augmentation to the input data and labels
        if args.dropout > 0:
            X = add_dropout(X, args.dropout, len(le_species.classes_))  # Add dropout to the input data
        if args.noise:
            X = add_noise(X, len(le_species.classes_))  # Add noise to the input data
        if args.log:
            X = add_log(X, len(le_species.classes_))  # Apply logarithmic transformation to the input data
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X = remove_features(args, X, features)  # Remove specific features from the input data
        if args.standardization:
            X, scaler = add_standardization(X, scaler=None)  # Add standardization to the input data and save the scaler
            set_scaler(args, scaler)
        train_dataloader = add_dataloader(args, X, y, shuffle=True)
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))
        model = MLP(**model_parameters)
        if args.criterion == 'imbal-noised-top-k':
            LastLinearLayer = NormedLinear
            model.fc2 = LastLinearLayer(model.fc2.in_features, len(le_header.classes_))
        device = get_device(args)  # Set the device to run the model on to be GPU
        if args.data_parallelism:
            model = torch.nn.DataParallel(model)  # Run the model parallelly
        model.to(device)  # Use GPU for model
        criterion = get_criterion(args, y)  # Instantiate loss class
        criterion.to(device)
        optimizer = get_optimizer(args, model)  # Instantiate optimizer class
        scheduler = get_scheduler(args, optimizer)  # Instantiate step learning scheduler class
        best_accuracy = -np.inf
        best_epoch = -1
        best_model_state = model.state_dict()  # Set up a variable to store the best model's state
        no_improvement_epochs = 0  # Set up a counter for the number of epochs without improvement
        start_training = time.time()
        for epoch in range(args.num_epochs):
            epoch_accuracy, model = self.train_epoch(args, model, train_dataloader, epoch, device, optimizer, criterion, le_header)  # Train the model for one epoch
            scheduler.step(epoch_accuracy)  # Decay Learning Rate, pass validation accuracy for tracking at every epoch
            if epoch_accuracy > best_accuracy:  
                best_accuracy = epoch_accuracy  
                best_epoch = epoch
                best_model_state = model.state_dict()  # Update the best model state
                set_mlp_model(args, best_model_state)  # Save the model
                print("\nSuccessfully saved model at Models/MLP.pth")
                no_improvement_epochs = 0
            else:  # Check if the accuracy is not the best seen so far
                no_improvement_epochs += 1  # Increment the number of epochs without improvement
                print("\nEpoch {} did not improve the model.".format(epoch))
            if no_improvement_epochs == args.num_iter_no_change:  # Check if the model's accuracy on the test set has not improved for the num_iter_no_change number of epochs               
                print(f'\nEarly stopping occurred at epoch {epoch} with best epoch = {best_epoch} and best {args.metric} = {best_accuracy:.4f}%.')
                break  # Stop the training process
        end_training = time.time()
        time_training = end_training - start_training
        if no_improvement_epochs != args.num_iter_no_change:
            print(f'\nEarly stopping dit not occur as best epoch = {best_epoch} and best {args.metric} = {best_accuracy:.4f}%.')
        print(f"Total time of the training: {time_training:.2f}s.")

    def predict_model(self, args, X, y, le_species, le_header):
        """
        Predict the habitat labels for given input data using a trained model.

        Args:
            args (argparse.Namespace): Command-line arguments.
            X (scipy.sparse._csr.csr_matrix): Input data features.
            y (numpy.ndarray): Input data labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): LabelEncoder for species classes.
            le_header (sklearn.preprocessing._label.LabelEncoder): LabelEncoder for habitat classes.

        Returns:
            None
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank features to input data
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features to input data
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features to input data
        if args.log:
            X = add_log(X, len(le_species.classes_))  # Add logarithmic features to input data
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X = remove_features(args, X, features)  # Remove specified features from input data
        if args.standardization:
            scaler = get_scaler(args)
            X, _ = add_standardization(X, scaler=scaler)  # Add standardized features to input data using the scaler
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Get model parameters
        model = MLP(**model_parameters)  # Create an instance of the MLP model
        if args.criterion == 'imbal-noised-top-k':
            LastLinearLayer = NormedLinear
            model.fc2 = LastLinearLayer(model.fc2.in_features, len(le_header.classes_))  # Replace the last linear layer with a custom NormedLinear layer
        device = get_device(args)  # Set the device to run the model on to be GPU
        if args.data_parallelism:
            model = torch.nn.DataParallel(model)  # Run the model in parallel if specified
        model = get_mlp_model(args, model)  # Load the trained model parameters
        model.to(device)  # Move the model to the specified device (GPU or CPU)
        X = torch.from_numpy(X.toarray()).float()  # Convert input data to a PyTorch tensor
        model.eval()  # Set the model to evaluation mode
        start_prediction = time.time()  # Start timing the prediction process
        with torch.no_grad():  # Turn off gradient computation to save memory
            predictions = model(X)  # Pass the input data through the model to obtain predictions
        end_prediction = time.time()  # End timing the prediction process
        time_prediction = end_prediction - start_prediction  # Compute the time taken for prediction
        predictions = scipy.special.softmax(predictions.cpu(), axis=1)  # Apply softmax function to convert logits to probabilities
        predictions = [np.argmax(sub_arr) for sub_arr in predictions]  # Get the indices of the maximum probability for each prediction
        predictions = add_predictions_decoding(predictions, le_header)  # Decode the predicted labels using the LabelEncoder for habitats
        set_predictions(args, predictions)  # Store the predictions in the specified location
        print("\nSuccessfully saved predictions at Data/predictions.txt")
        print(f"Total time of the prediction: {time_prediction:.2f}s.")  # Print the total time taken for prediction

    def evaluate_fold(self, args, X, y, le_species, le_header, fold, split_assignments):
        """
        Evaluate a single fold of the dataset for Habitat Distribution Modeling.

        Args:
            args (argparse.Namespace): Command-line arguments.
            X (scipy.sparse._csr.csr_matrix): Input data features.
            y (numpy.ndarray): Input data labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): LabelEncoder for species classes.
            le_header (sklearn.preprocessing._label.LabelEncoder): LabelEncoder for habitat classes.
            fold (int): Fold number for evaluation.
            split_assignments (numpy.ndarray): Array indicating the fold assignments for each sample.

        Returns:
            float: Best accuracy achieved in the fold.
        """
        print('\n' + '*'*11)
        print('* Fold: {} *'.format(fold))  # Print fold number
        print('*'*11 + '\n' + '-'*30)
        X_train, X_test, y_train, y_test = add_fold_assignments(X, y, split_assignments, fold)  # Split data into training and test sets based on fold assignments
        if args.endangered:
            eunis_red_list_crosswalks = get_eunis_red_list_crosswalks(args)
            endangered_red_list_habitats = get_endangered_red_list_habitats(args)
            endangered_eunis_habitats = get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header)  # Get endangered habitats
            X_test, y_test = add_endangered_habitats(X_test, y_test, endangered_eunis_habitats)   # Add endangered habitats to the test set
        if args.augmentation:
            if not args.endangered:
                labels = None
            else:
                labels = endangered_eunis_habitats
            X_train, y_train = add_augmentation(X_train, y_train, labels)  # Augment the training set with additional samples
        if args.dropout > 0:
            X_test = add_dropout(X_test, args.dropout, len(le_species.classes_))  # Apply dropout to the test set
        if args.noise:
            X_test = add_noise(X_test, len(le_species.classes_))  # Add noise to the test set
        if args.log:
            X_train = add_log(X_train, len(le_species.classes_))  # Apply logarithmic transformation to the training set
            X_test = add_log(X_test, len(le_species.classes_))  # Apply logarithmic transformation to the test set
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X_train = remove_features(args, X_train, features)  # Remove specific features from the training set
            X_test = remove_features(args, X_test, features)  # Remove specific features from the test set
        if args.standardization:
            X_train, scaler = add_standardization(X_train, scaler=None)  # Apply standardization to the training set
            X_test, scaler = add_standardization(X_test, scaler=scaler)  # Apply standardization to the test set using the same scaler
        train_dataloader = add_dataloader(args, X_train, y_train, shuffle=True)  # Create a dataloader for the training set
        test_dataloader = add_dataloader(args, X_test, y_test, shuffle=False)  # Create a dataloader for the test set
        model_parameters = get_model_parameters(args, X_train.shape[1], len(le_header.classes_))  # Get model parameters
        model = MLP(**model_parameters)  # Create an instance of the MLP model
        if args.criterion == 'imbal-noised-top-k':
            LastLinearLayer = NormedLinear
            model.fc2 = LastLinearLayer(model.fc2.in_features, len(le_header.classes_))  # Replace the last linear layer with a customized linear layer
        device = get_device(args)  # Set the device to run the model on (GPU or CPU)
        if args.data_parallelism:
            model = torch.nn.DataParallel(model)  # Use data parallelism to run the model on multiple GPUs
        model.to(device)  # Move the model to the specified device (GPU or CPU)
        criterion = get_criterion(args, y_train)  # Instantiate the loss function
        criterion.to(device)  # Move the loss function to the specified device (GPU or CPU)
        optimizer = get_optimizer(args, model)  # Instantiate the optimizer
        scheduler = get_scheduler(args, optimizer)  # Instantiate the learning rate scheduler
        best_fold_accuracy = -np.inf  # Initialize the best accuracy variable
        best_epoch = -1  # Initialize the best epoch variable
        no_improvement_epochs = 0  # Initialize the counter for epochs without improvement
        start_fold = time.time()  # Start timing the evaluation of the fold
        for epoch in range(args.num_epochs):
            accuracy, model = self.evaluate_epoch(args, model, train_dataloader, test_dataloader, epoch, device, optimizer, criterion, le_header)  # Evaluate the epoch and obtain the accuracy
            scheduler.step(accuracy)  # Decay Learning Rate, pass validation accuracy for tracking at every epoch
            if accuracy > best_fold_accuracy:  # Check if the model's state has the highest accuracy on the test set so far
                best_fold_accuracy = accuracy  # Update the best accuracy
                best_epoch = epoch  # Update the best epoch
                no_improvement_epochs = 0  # Reset the number of epochs without improvement
            else:  # Check if the test accuracy is not the best seen so far
                no_improvement_epochs += 1  # Increment the number of epochs without improvement
            if no_improvement_epochs == args.num_iter_no_change:  # Check if the model's accuracy on the test set has not improved for the num_iter_no_change number of epochs               
                print(f'\nEarly stopping occurred at epoch {epoch} with best epoch = {best_epoch} and best {args.metric} = {best_fold_accuracy}')
                break  # Stop the training process
        end_fold = time.time()  # End timing the evaluation of the fold
        time_fold = end_fold - start_fold  # Calculate the time taken for the evaluation of the fold
        if no_improvement_epochs != args.num_iter_no_change:
            print(f'\nEarly stopping dit not occur as best epoch = {best_epoch} and best {args.metric} = {best_fold_accuracy}')
        print(f"Total time for the evaluation of the fold: {time_fold:.2f}s.")
        return best_fold_accuracy

    def train_epoch(self, args, model, train_dataloader, epoch, device, optimizer, criterion, le_header):
        """
        Train the model for a single epoch.

        Args:
            args (argparse.Namespace): Arguments passed to the training process.
            model (models.mlp.MLP | torch.nn.parallel.data_parallel.DataParallel): The deep learning model.
            train_dataloader (torch.utils.data.dataloader.DataLoader): Dataloader for the training dataset.
            epoch (int): Current epoch number.
            device (torch.device): Device to run the model on (GPU or CPU).
            optimizer (torch.optim.sgd.SGD | torch.optim.adam.Adam | torch.optim.adamw.AdamW): Optimizer for updating model parameters.
            criterion (torch.nn.modules.loss.CrossEntropyLoss | pytopk.noised_losses.BalNoisedTopK | pytopk.noised_losses.ImbalNoisedTopK): Loss function.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.

        Returns:
            torch.Tensor: Accuracy achieved in the epoch.
            MLP: Updated model after training.
        """
        start_epoch = time.time()  # Start timing the epoch
        for batch, labels in train_dataloader:  # Iterate through train dataset
            batch = batch.requires_grad_().to(device)  # Load batches with gradient accumulation capabilities
            labels = labels.to(device)  # Use GPU for tensors
            optimizer.zero_grad()  # Clear gradients w.r.t. parameters
            outputs = model(batch)  # Forward pass to get output/logits
            loss = criterion(outputs, labels)  # Calculate Loss: softmax --> cross entropy loss
            loss.backward()  # Getting gradients w.r.t. parameters
            optimizer.step()  # Updating parameters
        end_epoch = time.time()  # End timing the epoch
        time_epoch = end_epoch - start_epoch  # Calculate the time taken for the epoch
        metric = get_metric(args, len(le_header.classes_))  # Instantiate the accuracy metric
        metric.to(device)  # Move the metric to the device
        model.eval()  # Set the model to evaluation mode
        with torch.no_grad():  # Disable gradient computation
            for batch, labels in train_dataloader:  # Iterate through the batches in the training dataset
                batch = batch.to(device)  # Move the batch to the device
                labels = labels.to(device)  # Move the labels to the device
                outputs = model(batch)  # Forward pass to get output/logits
                accuracy = metric(outputs, labels)  # Calculate the accuracy
        accuracy_epoch = metric.compute() * 100  # Compute the overall accuracy for the epoch
        print('\n' + '-'*30)
        print('Epoch {} completed.'.format(epoch))  # Print the epoch number
        print('LR: {}.'.format(optimizer.param_groups[0]['lr']))  # Print Learning Rate
        print(f"Time: {time_epoch:.2f}s.")  # Print Time
        print(f'Accuracy: {accuracy_epoch   :.4f}%.')  # Print Accuracy
        print('-'*30)
        model.train()  # Set the model back to training mode
        return accuracy_epoch, model

    def evaluate_epoch(self, args, model, train_dataloader, test_dataloader, epoch, device, optimizer, criterion, le_header):
        """
        Evaluate a single epoch of the training process.

        Args:
            args (argparse.Namespace): Training arguments.
            model (models.mlp.MLP | torch.nn.parallel.data_parallel.DataParallel): The deep learning model.
            train_dataloader (torch.utils.data.dataloader.DataLoader): Dataloader for the training dataset.
            test_dataloader (torch.utils.data.dataloader.DataLoader): Dataloader for the test dataset.
            epoch (int): Current epoch number.
            device (torch.device): Device to perform computations on (e.g., CPU or GPU).
            optimizer (torch.optim.sgd.SGD | torch.optim.adam.Adam | torch.optim.adamw.AdamW): Optimizer for updating model parameters.
            criterion (torch.nn.modules.loss.CrossEntropyLoss | pytopk.noised_losses.BalNoisedTopK | pytopk.noised_losses.ImbalNoisedTopK): Loss function criterion.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for the header.

        Returns:
            tuple: A tuple containing the accuracy of the epoch and the updated model.
        """
        start_epoch = time.time()  # Start timing the epoch
        for batch, labels in train_dataloader:  # Iterate through train dataset
            batch = batch.requires_grad_().to(device)  # Load batches with gradient accumulation capabilities
            labels = labels.to(device)  # Use GPU for tensors
            optimizer.zero_grad()  # Clear gradients w.r.t. parameters
            outputs = model(batch)  # Forward pass to get output/logits
            loss = criterion(outputs, labels)  # Calculate Loss: softmax --> cross entropy loss
            loss.backward()  # Getting gradients w.r.t. parameters
            optimizer.step()  # Updating parameters
        end_epoch = time.time()  # End timing the epoch
        time_epoch = end_epoch - start_epoch  # Calculate the time taken for the epoch
        metric = get_metric(args, len(le_header.classes_))  # Instantiate the accuracy metric
        metric.to(device)  # Move the metric to the device
        model.eval()  # Set the model to evaluation mode
        with torch.no_grad():  # Disable gradient computation
            for batch, labels in test_dataloader:  # Iterate through the batches in the test dataset
                batch = batch.to(device)  # Move the batch to the device
                labels = labels.to(device)  # Move the labels to the device
                outputs = model(batch)  # Forward pass to get output/logits
                accuracy = metric(outputs, labels)  # Calculate the accuracy
        accuracy_epoch = metric.compute() * 100  # Compute the overall accuracy for the epoch
        print('Epoch {} completed.'.format(epoch))  # Print Epoch
        print('LR: {}.'.format(optimizer.param_groups[0]['lr']))  # Print Learning Rate
        print(f"Time: {time_epoch:.2f}s.")  # Print Time
        print(f'Accuracy: {accuracy_epoch:.4f}%.')  # Print Accuracy
        print('-'*30)
        model.train()  # Set the model back to training mode
        return accuracy_epoch, model

    def run(self, args, X, y, le_species, le_header):
        """
        Run the MLP model based on the specified pipeline.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.pipeline == 'evaluation':
            self.evaluate_model(args, X, y, le_species, le_header)  # Run evaluation pipeline
        elif args.pipeline == 'training':
            self.train_model(args, X, y, le_species, le_header)  # Run training pipeline
        elif args.pipeline == 'prediction':
            self.predict_model(args, X, y, le_species, le_header)  # Run prediction pipeline

models/mlp.py
cesar-leblanc/hdm-framework
import sklearn.ensemble
import torch
import numpy as np
import time

from utils import get_metric, get_model_parameters

from data.load_data import get_split_assignments, get_eunis_red_list_crosswalks, get_endangered_red_list_habitats, get_endangered_eunis_habitats, get_rfc_model, get_scaler
from data.preprocess_data import add_fold_assignments, add_standardization, add_normalization, add_dropout, add_binarization, add_log, add_noise, add_rank, remove_features, add_augmentation, add_endangered_habitats, add_predictions_decoding
from data.save_data import set_scaler, set_rfc_model, set_predictions

class RFC(sklearn.ensemble.RandomForestClassifier):
    """
    Random Forest Classifier (RFC) model for habitat distribution modeling.
    """
    def __init__(self, n_estimators=200, criterion='gini', max_depth=None, min_samples_split=2, min_samples_leaf=1, min_weight_fraction_leaf=0.0, max_features='sqrt', max_leaf_nodes=None, min_impurity_decrease=0.0, bootstrap=True, oob_score=False, n_jobs=-1, random_state=123, verbose='n_estimators', warm_start=False, class_weight=None, ccp_alpha=0.0, max_samples=None):
        """
        Initialize the Random Forest Classifier.

        Args:
            n_estimators (int, optional): The number of trees in the forest. Defaults to 200.
            criterion (str, optional): The function to measure the quality of a split. Defaults to 'gini'.
            max_depth (int, optional): The maximum depth of the tree. Defaults to None.
            min_samples_split (int, optional): The minimum number of samples required to split an internal node. Defaults to 2.
            min_samples_leaf (int, optional): The minimum number of samples required to be at a leaf node. Defaults to 1.
            min_weight_fraction_leaf (float, optional): The minimum weighted fraction of the sum total of weights (of all the input samples) required to be at a leaf node. Defaults to 0.0.
            max_features (int | float, optional): The number of features to consider when looking for the best split. Defaults to 'sqrt'.
            max_leaf_nodes (int, optional): The value for which trees grow in a best-first fashion. Defaults to None.
            min_impurity_decrease (float, optional): The value for which a node will be split if this split induces a decrease of the impurity greater than or equal to this value. Defaults to 0.0.
            bootstrap (bool, optional): The choice to use bootstrap samples when building trees. Defaults to True.
            oob_score (bool, optional): The choice to use out-of-bag samples to estimate the generalization score. Defaults to False.
            n_jobs (int | None, optional): The number of jobs to run in parallel. Defaults to -1.
            random_state (int | RandomState, optional): The value that controls both the randomness of the bootstrapping of the samples used when building trees and the sampling of the features to consider when looking for the best split at each node. Defaults to 123.
            verbose (int, optional): The value that controls the verbosity when fitting and predicting. Defaults to 'n_estimators'.
            warm_start (bool, optional): The choice to reuse the solution of the previous call to fit and add more estimators to the ensemble or to just fit a whole new forest. Defaults to False.
            class_weight (str | dict | list | None, optional): The weights associated with classes. Defaults to None.
            ccp_alpha (float, optional): The complexity parameter used for Minimal Cost-Complexity Pruning. Defaults to 0.0.
            max_samples (int | float, optional): The number of samples to draw from X to train each base estimator if bootstrap is used. Defaults to None.
        """
        super(RFC, self).__init__()  # Initialize the parent class, `sklearn.ensemble.RandomForestClassifier`
        self.n_estimators = n_estimators  # The number of trees in the forest
        self.criterion = criterion  # The function to measure the quality of a split
        self.max_depth = max_depth  # The maximum depth of the tree
        self.min_samples_split = min_samples_split  # The minimum number of samples required to split an internal node
        self.min_samples_leaf = min_samples_leaf  # The minimum number of samples required to be at a leaf node
        self.min_weight_fraction_leaf = min_weight_fraction_leaf  # The minimum weighted fraction of the sum total of weights (of all the input samples) required to be at a leaf node
        self.max_features = max_features  # The number of features to consider when looking for the best split
        self.max_leaf_nodes = max_leaf_nodes  # The value for which trees grow in best-first fashion
        self.min_impurity_decrease = min_impurity_decrease  # The value for which a node will be split if this split induces a decrease of the impurity greater than or equal to this value
        self.bootstrap = bootstrap  # The choice to use bootstrap samples when building trees
        self.oob_score = oob_score  # The choice to use out-of-bag samples to estimate the generalization score
        self.n_jobs = n_jobs  # The number of jobs to run in parallel
        self.random_state = random_state  # The value that controls both the randomness of the bootstrapping of the samples used when building trees and the sampling of the features to consider when looking for the best split at each node
        self.verbose = verbose  # The value that controls the verbosity when fitting and predicting
        self.warm_start = warm_start  # The choice to reuse the solution of the previous call to fit and add mor estimators to the ensemble or to just fit a whole new forest
        self.class_weight = class_weight  # The weights associated with classes
        self.ccp_alpha = ccp_alpha  # The complexity parameter used for Minimal Cost-Complexity Pruning
        self.max_samples = max_samples  # The number of samples to draw from X to train each base estimator if bootstrap is used
        
    def evaluate_model(self, args, X, y, le_species, le_header):
        """
        Evaluate the model's performance using the specified evaluation settings.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank as a feature
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features
        std_acc = []
        avg_acc = 0.0  # Initialize variables to track the average accuracy
        n_samples = 0  # Initialize variables to track the number of samples
        split_assignments = get_split_assignments(args)  # Get split assignments for cross-validation
        best_accuracy = -np.inf  # Initialize variable to track the best accuracy
        best_fold = -1  # Initialize variable to track the best fold
        start_evaluation = time.time()  # Start time of evaluation
        for fold in range(args.n_folds):
            best_fold_accuracy = self.evaluate_fold(args, X, y, le_species, le_header, fold, split_assignments)  # Perform evaluation for a fold
            if best_fold_accuracy > best_accuracy:  # Check if current fold improves the model
                best_accuracy = best_fold_accuracy  # Update the best accuracy
                best_fold = fold  # Update the best fold
                print("\nFold {} did improve the model.".format(best_fold))
            else:
                print("\nFold {} did not improve the model.".format(fold))
            std_acc.append(best_fold_accuracy)  # Add best fold accuracy to the list
            avg_acc += best_fold_accuracy  # Accumulate the accuracy
            n_samples += 1  # Accumulate the number of samples
        end_evaluation = time.time()  # End time of evaluation
        time_evaluation = end_evaluation - start_evaluation  # Total evaluation time
        avg_acc /= n_samples  # Compute the average accuracy
        std_acc = torch.std(torch.FloatTensor(std_acc))  # Compute the standard deviation of accuracies
        print(f"\nMean {args.metric} and standard deviation after {args.n_folds} folds of {args.n_estimators} estimators: {avg_acc:.4f}% & {std_acc:.4f}.")
        print(f"Total time of the evaluation: {time_evaluation:.2f}s.")

    def train_model(self, args, X, y, le_species, le_header):
        """
        Train the model using the specified training settings.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank as a feature
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features
        if args.endangered:
            eunis_red_list_crosswalks = get_eunis_red_list_crosswalks(args)  # Get EUNIS-Red list crosswalks
            endangered_red_list_habitats = get_endangered_red_list_habitats(args)  # Get endangered Red list habitats
            endangered_eunis_habitats = get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header)  # Get endangered EUNIS habitats
            X, y = add_endangered_habitats(X, y, endangered_eunis_habitats)  # Add endangered habitats to the data
        if args.augmentation:
            if not args.endangered:
                labels = None
            else:
                labels = endangered_eunis_habitats
            X, y = add_augmentation(X, y, labels)  # Perform data augmentation
        if args.dropout > 0:
            X = add_dropout(X, args.dropout, len(le_species.classes_))  # Add dropout regularization
        if args.noise:
            X = add_noise(X, len(le_species.classes_))  # Add noise to the features
        if args.log:
            X = add_log(X, len(le_species.classes_))  # Apply logarithm to the feature
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X = remove_features(args, X, features)  # Remove specified features
        if args.standardization:
            X, scaler = add_standardization(X, scaler=None)  # Standardize the features
            set_scaler(args, scaler)  # Set the scaler for future use
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Get model parameters
        model = RFC(**model_parameters)  # Create an instance of the Random Forest Classifier
        print()
        start_training = time.time()  # Start time of training
        model.fit(X, y)  # Train the model
        end_training = time.time()  # End time of training
        time_training = end_training - start_training  # Total training time
        print(f"\nTotal time of the training: {time_training:.2f}s.\n")
        y_probas = model.predict_proba(X)  # Get predicted probabilities
        metric = get_metric(args, len(le_header.classes_))  # Get the evaluation metric
        accuracy = metric(torch.from_numpy(y_probas), torch.from_numpy(y)) * 100  # Compute accuracy using the metric
        print(f"\nMean {args.metric} with {args.n_estimators} estimators: {accuracy:.4f}%.")
        set_rfc_model(args, model)  # Set the trained model for future use
        print("\nSuccessfully saved model at Models/RFC.joblib")
        
    def predict_model(self, args, X, y, le_species, le_header):
        """
        Use the trained model to make predictions on the input data.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank as a feature
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features
        if args.log:
            X = add_log(X, len(le_species.classes_))  # Apply logarithm to the features
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X = remove_features(args, X, features)  # Remove specified features
        if args.standardization:
            scaler = get_scaler(args)  # Get the scaler used during training
            X, _ = add_standardization(X, scaler=scaler)  # Standardize the features using the scaler
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Get model parameters
        model = RFC(**model_parameters)  # Create an instance of the Random Forest Classifier
        model = get_rfc_model(args, model)  # Load the trained model
        print()
        start_prediction = time.time()  # Start time of prediction
        predictions = model.predict(X)  # Make predictions
        end_prediction = time.time()  # End time of prediction
        time_prediction = end_prediction - start_prediction  # Total prediction time
        predictions = add_predictions_decoding(predictions, le_header)  # Decode the predictions using label encoder
        set_predictions(args, predictions)  # Set the predictions for future use
        print("\nSuccessfully saved predictions at Data/predictions.txt")
        print(f"Total time of the prediction: {time_prediction:.2f}s.")
    
    def evaluate_fold(self, args, X, y, le_species, le_header, fold, split_assignments):
        """
        Evaluate the model performance on a specific fold of the data.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
            fold (int): Fold index.
            split_assignments (numpy.ndarray): Assignment of samples to folds.

        Returns:
            accuracy: Accuracy achieved on the fold.
        """
        print('\n' + '*'*11)
        print('* Fold: {} *'.format(fold))  # Print fold
        print('*'*11 + '\n')
        X_train, X_test, y_train, y_test = add_fold_assignments(X, y, split_assignments, fold)  # Split data into train and test sets
        if args.endangered:
            eunis_red_list_crosswalks = get_eunis_red_list_crosswalks(args)
            endangered_red_list_habitats = get_endangered_red_list_habitats(args)
            endangered_eunis_habitats = get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header)
            X_test, y_test = add_endangered_habitats(X_test, y_test, endangered_eunis_habitats)  # Add endangered habitats to the test set
        if args.augmentation:
            if not args.endangered:
                labels = None
            else:
                labels = endangered_eunis_habitats  
            X_train, y_train = add_augmentation(X_train, y_train, labels)  # Add data augmentation to the training set
        if args.dropout > 0:
            X_test = add_dropout(X_test, args.dropout, len(le_species.classes_))  # Apply dropout to the test set
        if args.noise:
            X_test = add_noise(X_test, len(le_species.classes_))  # Add noise to the test set
        if args.log:
            X_train = add_log(X_train, len(le_species.classes_))  # Apply logarithm to the training set
            X_test = add_log(X_test, len(le_species.classes_))  # Apply logarithm to the test set
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X_train = remove_features(args, X_train, features)  # Remove specified features from the training set
            X_test = remove_features(args, X_test, features)  # Remove specified features from the test set
        if args.standardization:
            X_train, scaler = add_standardization(X_train, scaler=None)  # Standardize the training set
            X_test, scaler = add_standardization(X_test, scaler=scaler)  # Standardize the test set using the same scaler
        model_parameters = get_model_parameters(args, X_train.shape[1], len(le_header.classes_))  # Get model parameters
        model = RFC(**model_parameters)  # Create an instance of the Random Forest Classifier
        start_fold = time.time()  # Start time of the fold
        model.fit(X_train, y_train)  # Train the model
        end_fold = time.time()  # End time of the fold
        time_fold = end_fold - start_fold  # Total fold time
        y_probas = model.predict_proba(X_test)  # Get predicted probabilities for the test set
        metric = get_metric(args, len(le_header.classes_))  # Get the evaluation metric
        accuracy = metric(torch.from_numpy(y_probas), torch.from_numpy(y_test)) * 100  # Calculate accuracy
        print(f"\nAccuracy: {accuracy:.4f}%.")  # Print Accuracy
        print(f"Time: {time_fold:.2f}s.")  # Print Time
        return accuracy

    def run(self, args, X, y, le_species, le_header):
        """
        Run the RFC model based on the specified pipeline.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.pipeline == 'evaluation':
            self.evaluate_model(args, X, y, le_species, le_header)  # Run evaluation pipeline
        elif args.pipeline == 'training':
            self.train_model(args, X, y, le_species, le_header)  # Run training pipeline
        elif args.pipeline == 'prediction':
            self.predict_model(args, X, y, le_species, le_header)  # Run prediction pipeline

models/rfc.py
cesar-leblanc/hdm-framework
import pytorch_tabnet.tab_model
import torch
import numpy as np
import time

from utils import get_metric, get_model_parameters, get_fit_parameters

from data.load_data import get_split_assignments, get_eunis_red_list_crosswalks, get_endangered_red_list_habitats, get_endangered_eunis_habitats, get_tnc_model, get_scaler
from data.preprocess_data import add_fold_assignments, add_standardization, add_normalization, add_dropout, add_binarization, add_log, add_noise, add_rank, remove_features, add_augmentation, add_endangered_habitats, add_predictions_decoding
from data.save_data import set_scaler, set_tnc_model, set_predictions

class TNC(pytorch_tabnet.tab_model.TabNetClassifier):
    """
    TabNet Classifier for habitat distribution modeling.
    """
    def __init__(self, n_d=16, n_a=16, n_steps=3, gamma=1.3, cat_idxs=[], cat_dims=[], cat_emb_dim=1, n_independent=1, n_shared=3, epsilon=1e-15, seed=123, momentum=0.02, clip_value=None, lambda_sparse=1e-3, optimizer_fn='adam', optimizer_params={'lr':0.02}, scheduler_fn='step', scheduler_params={"step_size": 10, "gamma": 0.9}, model_name='TabNetClassifier', verbose=1, device_name='auto', input_dim= None, output_dim=None, mask_type='sparsemax', grouped_features=[], n_shared_decoder=1, n_indep_decoder=1):
        """
        Initialize the TabNet Classifier.

        Args:
            n_d (int, optional): The width of the decision prediction layer. Defaults to 16.
            n_a (int, optional): The width of the attention embedding for each mask. Defaults to 16.
            n_steps (int, optional): The number of steps in the architecture (usually between 3 and 10). Defaults to 3.
            gamma (float, optional): The coefficient for feature reusage in the masks. Defaults to 1.3.
            cat_idxs (list, optional): The list of categorical features indices. Defaults to [].
            cat_dims (list, optional): The list of categorical features number of modalities (number of unique values for a categorical feature). Defaults to [].
            cat_emb_dim (list, optional): The list of embeddings size for each categorical feature. Defaults to 1.
            n_independent (int, optional): The number of independent Gated Linear Units layers at each step. Defaults to 1.
            n_shared (int, optional): The number of shared Gated Linear Units at each step. Defaults to 3.
            epsilon (float, optional): The value should be left untouched. Defaults to 1e-15.
            seed (int, optional): The random seed for reproducibility. Defaults to 123.
            momentum (float, optional): The momentum for batch normalization.Defaults to 0.02.
            clip_value (float | tuple, optional): The gradient will be clipped at the given value. Defaults to None.
            lambda_sparse (float, optional): The extra sparsity loss coefficient. Defaults to 1e-3.
            optimizer_fn (callable, optional): The PyTorch optimizer function. Defaults to 'adam'.
            optimizer_params (dict, optional): The parameters compatible with the PyTorch optimizer function used to initialize the optimizer. Defaults to {'lr':0.02}
            scheduler_fn (callable, optional): The PyTorch scheduler to change learning rates during training. Defaults to 'step'.
            scheduler_params (dict, optional): The dictionary of parameters to apply to the PyTorch scheduler. Defaults to {"step_size": 10, "gamma": 0.9}.
            model_name (str, optional): The name of the model used for saving in disk. Defaults to 'TabNetClassifier'.
            verbose (int, optional): The verbosity for notebooks plots. Defaults to 1.
            device_name (str, optional): The name of the device used for training. Defaults to 'auto'.
            input_dim (int, optional): The number of input features. Defaults to None.
            output_dim (int, optional): The number of output classes. Defaults to None.
            mask_type (str, optional): The masking function to use for selecting features. Defaults to 'sparsemax'.
            grouped_features (list): The groups to allow the model to share its attention across features inside the same group. Defaults to [].
            n_shared_decoder (int, optional): The number of shared GLU blocks in the decoder. Defaults to 1.
            n_indep_decoder (int, optional): The number of independent GLU blocks in the decoder. Defaults to 1.
        """
        super(TNC, self).__init__()  # Initialize the parent class, `pytorch_tabnet.tab_model.TabNetClassifier`
        self.n_d = n_d  # The width of the decision prediction layer
        self.n_a = n_a  # The width of the attention embedding for each mask
        self.n_steps = n_steps  # The number of steps in the architecture (usually between 3 and 10)
        self.gamma = gamma  # The coefficient for feature reusage in the masks
        self.cat_idxs = cat_idxs  # The list of categorical features indices
        self.cat_dims = cat_dims  # The list of categorical features number of modalities (number of unique values for a categorical feature)
        self.cat_emb_dim = cat_emb_dim  # The list of embeddings size for each categorical features
        self.n_independent = n_independent  # The number of independent Gated Linear Units layers at each step
        self.n_shared = n_shared  # The number of shared Gated Linear Units at each step
        self.epsilon = epsilon  # The value should be left untouched
        self.seed = seed  # The random seed for reproducibility
        self.momentum = momentum  # The momentum for batch normalization
        self.clip_value = clip_value  # The gradient will be clipped at the given value
        self.lambda_sparse = lambda_sparse  # The extra sparsity loss coefficient
        self.optimizer_fn = optimizer_fn  # The PyTorch optimizer function
        self.optimizer_params = optimizer_params  # The parameters compatible with the PyTorch optimizer function used to initialize the optimizer
        self.scheduler_fn = scheduler_fn  # The PyTorch scheduler to change learning rates during training
        self.scheduler_params = scheduler_params  # The dictionnary of parameters to apply to the PyTorch scheduler
        self.model_name = model_name  # The name of the model used for saving in disk
        self.verbose = verbose  # The verbosity for notebooks plots
        self.device_name = device_name  # The name of the device used for training
        self.input_dim = input_dim  # The number of input features
        self.output_dim = output_dim  # The number of output classes
        self.mask_type = mask_type  # The masking function to use for selection features
        self.grouped_features = grouped_features  # The groups to allow the model to share its attention across features inside a same group
        self.n_shared_decoder = n_shared_decoder  # The number of shared GLU block in decoder
        self.n_indep_decoder = n_indep_decoder  # The number of independent GLU block in decoder

    def evaluate_model(self, args, X, y, le_species, le_header):
        """
        Evaluate the model's performance using the specified evaluation settings.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank as a feature
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features
        std_acc = []
        avg_acc = 0.0  # Initialize variables to track the average accuracy
        n_samples = 0  # Initialize variables to track the number of samples
        split_assignments = get_split_assignments(args)  # Get split assignments for cross-validation
        best_accuracy = -np.inf  # Initialize variable to track the best accuracy
        best_fold = -1  # Initialize variable to track the best fold
        start_evaluation = time.time()  # Start time of evaluation
        for fold in range(args.n_folds):
            best_fold_accuracy = self.evaluate_fold(args, X, y, le_species, le_header, fold, split_assignments)  # Perform evaluation for a fold
            if best_fold_accuracy > best_accuracy:  # Check if current fold improves the model
                best_accuracy = best_fold_accuracy  # Update the best accuracy
                best_fold = fold  # Update the best fold
                print("\nFold {} did improve the model.".format(best_fold))
            else:
                print("\nFold {} did not improve the model.".format(fold))
            std_acc.append(best_fold_accuracy)  # Add best fold accuracy to the list
            avg_acc += best_fold_accuracy  # Accumulate the accuracy
            n_samples += 1  # Accumulate the number of samples
        end_evaluation = time.time()  # End time of evaluation
        time_evaluation = end_evaluation - start_evaluation  # Total evaluation time
        avg_acc /= n_samples  # Compute the average accuracy
        std_acc = torch.std(torch.FloatTensor(std_acc))  # Compute the standard deviation of accuracies
        print(f"\nMean {args.metric} and standard deviation after {args.n_folds} folds of {args.max_epochs} epochs: {avg_acc:.4f}% & {std_acc:.4f}.")
        print(f"Total time of the evaluation: {time_evaluation:.2f}s.")
        
    def train_model(self, args, X, y, le_species, le_header):
        """
        Train the model using the specified training settings.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank as a feature
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features
        if args.endangered:
            eunis_red_list_crosswalks = get_eunis_red_list_crosswalks(args)  # Get EUNIS-Red list crosswalks
            endangered_red_list_habitats = get_endangered_red_list_habitats(args)  # Get endangered Red list habitats
            endangered_eunis_habitats = get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header)  # Get endangered EUNIS habitats
            X, y = add_endangered_habitats(X, y, endangered_eunis_habitats)  # Add endangered habitats to the data
        if args.augmentation:
            if not args.endangered:
                labels = None
            else:
                labels = endangered_eunis_habitats
            X, y = add_augmentation(X, y, labels)  # Perform data augmentation
        if args.dropout > 0:
            X = add_dropout(X, args.dropout, len(le_species.classes_))  # Add dropout regularization
        if args.noise:
            X = add_noise(X, len(le_species.classes_))  # Add noise to the features
        if args.log:
            X = add_log(X, len(le_species.classes_))  # Apply logarithm to the feature
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X = remove_features(args, X, features)  # Remove specified features
        if args.standardization:
            X, scaler = add_standardization(X, scaler=None)  # Standardize the features
            set_scaler(args, scaler)  # Set the scaler for future use
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Get model parameters
        model = TNC(**model_parameters)  # Create an instance of the TabNet Classifier
        fit_parameters = get_fit_parameters(args, X, X, y, y)  # Get fit parameters
        print()
        start_training = time.time()  # Start time of training
        model.fit(**fit_parameters)  # Train the model
        end_training = time.time()  # End time of training
        time_training = end_training - start_training  # Total training time
        print(f"\nTotal time of the training: {time_training:.2f}s.")
        y_probas = model.predict_proba(X)  # Get predicted probabilities
        metric = get_metric(args, len(le_header.classes_))  # Get the evaluation metric
        accuracy = metric(torch.from_numpy(y_probas), torch.from_numpy(y)) * 100  # Compute accuracy using the metric
        print(f"\nMean {args.metric} with {args.n_steps} steps: {accuracy:.4f}%.")
        set_tnc_model(args, model)  # Set the trained model for future use
        
    def predict_model(self, args, X, y, le_species, le_header):
        """
        Use the trained model to make predictions on the input data.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank as a feature
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Add binarized features
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Add normalized features
        if args.log:
            X = add_log(X, len(le_species.classes_))  # Apply logarithm to the features
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X = remove_features(args, X, features)  # Remove specified features
        if args.standardization:
            scaler = get_scaler(args)  # Get the scaler used during training
            X, _ = add_standardization(X, scaler=scaler)  # Standardize the features using the scaler
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Get model parameters
        model = TNC(**model_parameters)  # Create an instance of the TabNet Classifier
        model = get_tnc_model(args, model)  # Load the trained model
        start_prediction = time.time()  # Start time of prediction
        predictions = model.predict(X)  # Make predictions
        end_prediction = time.time()  # End time of prediction
        time_prediction = end_prediction - start_prediction  # Total prediction time
        predictions = add_predictions_decoding(predictions, le_header)  # Decode the predictions using label encoder
        set_predictions(args, predictions)  # Set the predictions for future use
        print("\nSuccessfully saved predictions at Data/predictions.txt")
        print(f"Total time of the prediction: {time_prediction:.2f}s.")
    
    def evaluate_fold(self, args, X, y, le_species, le_header, fold, split_assignments):
        """
        Evaluate the model performance on a specific fold of the data.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
            fold (int): Fold index.
            split_assignments (numpy.ndarray): Assignment of samples to folds.

        Returns:
            accuracy: Accuracy achieved on the fold.
        """
        print('\n' + '*'*11)
        print('* Fold: {} *'.format(fold))  # Print fold
        print('*'*11 + '\n')
        X_train, X_test, y_train, y_test = add_fold_assignments(X, y, split_assignments, fold)  # Split data into train and test sets
        if args.endangered:
            eunis_red_list_crosswalks = get_eunis_red_list_crosswalks(args)
            endangered_red_list_habitats = get_endangered_red_list_habitats(args)
            endangered_eunis_habitats = get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header)
            X_test, y_test = add_endangered_habitats(X_test, y_test, endangered_eunis_habitats)  # Add endangered habitats to the test set
        if args.augmentation:
            if not args.endangered:
                labels = None
            else:
                labels = endangered_eunis_habitats
            X_train, y_train = add_augmentation(X_train, y_train, labels)  # Add data augmentation to the training set
        if args.dropout > 0:
            X_test = add_dropout(X_test, args.dropout, len(le_species.classes_))  # Apply dropout to the test set
        if args.noise:
            X_test = add_noise(X_test, len(le_species.classes_))  # Add noise to the test set
        if args.log:
            X_train = add_log(X_train, len(le_species.classes_))  # Apply logarithm to the training set
            X_test = add_log(X_test, len(le_species.classes_))  # Apply logarithm to the test set
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X_train = remove_features(args, X_train, features)  # Remove specified features from the training set
            X_test = remove_features(args, X_test, features)  # Remove specified features from the test set
        if args.standardization:
            X_train, scaler = add_standardization(X_train, scaler=None)  # Standardize the training set
            X_test, scaler = add_standardization(X_test, scaler=scaler)  # Standardize the test set using the same scaler
        model_parameters = get_model_parameters(args, X_train.shape[1], len(le_header.classes_))  # Get model parameters
        model = TNC(**model_parameters)  # Create an instance of the TabNet Classifier
        fit_parameters = get_fit_parameters(args, X_train, X_test, y_train, y_test)  # Get fit parameters
        start_fold = time.time()  # Start time of the fold
        model.fit(**fit_parameters)  # Train the model
        end_fold = time.time()  # End time of the fold
        time_fold = end_fold - start_fold  # Total fold time
        y_probas = model.predict_proba(X_test)  # Get predicted probabilities for the test set
        metric = get_metric(args, len(le_header.classes_))  # Get the evaluation metric
        accuracy = metric(torch.from_numpy(y_probas), torch.from_numpy(y_test)) * 100  # Calculate accuracy
        print(f"\nAccuracy: {accuracy:.4f}%.")  # Print Accuracy
        print(f"Time: {time_fold:.2f}s.")  # Print Time
        return accuracy

    def run(self, args, X, y, le_species, le_header):
        """
        Run the TNC model based on the specified pipeline.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.pipeline == 'evaluation':
            self.evaluate_model(args, X, y, le_species, le_header)  # Run evaluation pipeline
        elif args.pipeline == 'training':
            self.train_model(args, X, y, le_species, le_header)  # Run training pipeline
        elif args.pipeline == 'prediction':
            self.predict_model(args, X, y, le_species, le_header)  # Run prediction pipeline

models/tnc.py
cesar-leblanc/hdm-framework
import xgboost
import torch
import numpy as np
import time

from utils import get_metric, get_model_parameters

from data.load_data import get_split_assignments, get_eunis_red_list_crosswalks, get_endangered_red_list_habitats, get_endangered_eunis_habitats, get_xgb_model, get_scaler
from data.preprocess_data import add_fold_assignments, add_standardization, add_normalization, add_dropout, add_binarization, add_log, add_noise, add_rank, remove_features, add_augmentation, add_endangered_habitats, add_predictions_decoding
from data.save_data import set_scaler, set_xgb_model, set_predictions

class XGB(xgboost.XGBClassifier):
    """
    XGBoost Classifier for habitat distribution modeling.
    """
    def __init__(self, n_estimators=200, max_depth=100, max_leaves=0, max_bin=256, grow_policy='depthwise', learning_rate=0.1, verbosity=1, objective='multi:softmax', booster='gbtree', tree_method='auto', n_jobs=-1, gamma=0, min_child_weight=1, max_delta_step=0, subsample=0.8, sampling_method='uniform', colsample_bytree=1, colsample_bylevel=1, colsample_bynode=1, reg_alpha=0, reg_lambda=1, scale_pos_weight=None, base_score=None, random_state=123, missing=float('nan'), num_parallel_tree=1, monotone_constraints=None, interaction_constraints=None, importance_type=None, gpu_id=None, validate_parameters=None, predictor='auto', enable_categorical=None, feature_types=None, max_cat_to_onehot=None, max_cat_threshold=None, eval_metric='merror', early_stopping_rounds=10, callbacks=None, use_label_encoder=None):
        """
        Initialize the XGBoost Classifier.

        Args:
            n_estimators (int, optional): The number of boosting rounds. Defaults to 200.
            max_depth (int, optional): The maximum tree depth for base learners. Defaults to 100.
            max_leaves (int, optional): The maximum number of leaves; 0 indicates no limit. Defaults to 0.
            max_bin (int, optional): The maximum number of bins per feature if using histogram-based algorithm. Defaults to 256.
            grow_policy (str, optional): The tree growing policy. Defaults to 'depthwise'.
            learning_rate (float, optional): The boosting learning rate (xgb's 'eta'). Defaults to 0.1.
            verbosity (int, optional): The degree of verbosity. Defaults to 1.
            objective (str | callable, optional): The learning task and the corresponding learning objective or a custom objective function to be used. Defaults to 'multi:softmax'.
            booster (str, optional): The booster to use. Defaults to 'gbtree'.
            tree_method (str, optional): The tree method to use. Defaults to 'auto'.
            n_jobs (int, optional): The number of parallel threads used to run xgboost. Defaults to -1.
            gamma (float, optional): The minimum loss reduction required to make a further partition on a leaf node of the tree. Defaults to 0.
            min_child_weight (float, optional): The minimum sum of instance weight (hessian) needed in a child. Defaults to 1.
            max_delta_step (float, optional): The maximum delta step we allow each tree's weight estimation to be. Defaults to 0.
            subsample (float, optional): The subsample ratio of the training instance. Defaults to 0.8.
            sampling_method (str, optional): The sampling method. Defaults to 'uniform'.
            colsample_bytree (float, optional): The subsample ratio of columns when constructing each tree. Defaults to 1.
            colsample_bylevel (float, optional): The subsample ratio of columns for each level. Defaults to 1.
            colsample_bynode (float, optional): The subsample ratio of columns for each split. Defaults to 1.
            reg_alpha (float, optional): The L1 regularization term on weights (xgb's alpha). Defaults to 0.
            reg_lambda (float, optional): The L2 regularization term on weights (xgb's lambda). Defaults to 1.
            scale_pos_weight (float, optional): The balancing of positive and negative weights. Defaults to None.
            base_score (float, optional): The initial prediction score of all instances, global bias. Defaults to None.
            random_state (int, optional): The random number seed. Defaults to 123.
            missing (float | None, optional): The value in the data which needs to be present as a missing value. Defaults to float('nan').
            num_parallel_tree (int, optional): The value used for boosting random forest. Defaults to 1.
            monotone_constraints (str | list, optional): The constraint of variable monotonicity. Defaults to None.
            interaction_constraints (str | list, optional): The constraints for interaction representing permitted interactions. Defaults to None.
            importance_type (str, optional): The feature importance type for the feature importances property. Defaults to None.
            gpu_id (int, optional): The device ordinal. Defaults to None.
            validate_parameters (bool, optional): The choice to give warnings for unknown parameter. Defaults to None.
            predictor (str, optional): The value forces XGBoost to use a specific predictor. Defaults to 'auto'.
            enable_categorical (bool, optional): The support for categorical data. Defaults to None.
            feature_types (str | list, optional): The value used for specifying feature types without constructing a dataframe. Defaults to None.
            max_cat_to_onehot (int, optional): The threshold for deciding whether XGBoost should use one-hot encoding based split for categorical data. Defaults to None.
            max_cat_threshold (int, optional): The maximum number of categories considered for each split. Defaults to None.
            eval_metric (str, optional): The metric used for monitoring the training result and early stopping. Defaults to 'merror'.
            early_stopping_rounds (int, optional): The early stopping activation. Defaults to 10.
            callbacks (list, optional): The list of callback functions that are applied at end of each iteration. Defaults to None.
            use_label_encoder (bool, optional): The choice to use label encoder. Defaults to None.
        """
        super(XGB, self).__init__()  # Initialize the parent class, `xgboost.XGBClassifier`
        self.n_estimators = n_estimators  # The number of boosting rounds
        self.max_depth = max_depth  # The maximum tree depth for base learners
        self.max_leaves = max_leaves  # The maximum number of leaves; 0 indicates no limit
        self.max_bin = max_bin  # The maximum number of bins per feature if using histogram-based algorithm
        self.grow_policy = grow_policy  # The tree growing policy
        self.learning_rate = learning_rate  # The boosting learning rate (xgb's 'eta')
        self.verbosity = verbosity  # The degree of verbosity
        self.objective = objective  # The learning task and the corresponding learning objective or a custom objective function to be used
        self.booster = booster  # The booster to use
        self.tree_method = tree_method  # The tree method to use
        self.n_jobs = n_jobs  # The number of parallel threads used to run xgboost
        self.gamma = gamma  # The minimum loss reduction required to make a further partition on a leaf node of the tree
        self.min_child_weight = min_child_weight  # The minimum sum of instance weight (hessian) needed in a child
        self.max_delta_step = max_delta_step  # The maximum delta step we allow each tree's weight estimation to be
        self.subsample = subsample  # The subsample ratio of the training instance
        self.sampling_method = sampling_method  # The sampling method
        self.colsample_bytree = colsample_bytree  # The subsample ratio of columns when constructing each tree
        self.colsample_bylevel = colsample_bylevel  # The subsample ratio of columns for each level
        self.colsample_bynode = colsample_bynode  # The subsample ratio of columns for each split
        self.reg_alpha = reg_alpha  # The L1 regularization term on weights (xgb's alpha)
        self.reg_lambda = reg_lambda  # The L2 regularization term on weights (xgb's lambda)
        self.scale_pos_weight = scale_pos_weight  # The balancing of positive and negative weights
        self.base_score = base_score  # The initial prediction score of all instances, global bias
        self.random_state = random_state  # The random number seed
        self.missing = missing  # The value in the data which needs to be present as a missing value
        self.num_parallel_tree = num_parallel_tree  # The value used for boosting random forest
        self.monotone_constraints = monotone_constraints  # The constraint of variable monotonicity
        self.interaction_constraints = interaction_constraints  # The constraints for interaction representing permitted interactions
        self.importance_type = importance_type  # The feature importance type for the feature importances property
        self.gpu_id = gpu_id  # The device ordinal
        self.validate_parameters = validate_parameters  # The choice to give warnings for unknown parameter
        self.predictor = predictor  # The value forces XGBoost to use specific predictor
        self.enable_categorical = enable_categorical  # The support for categorical data
        self.feature_types = feature_types  # The value used for specifying feature types without constructing a dataframe
        self.max_cat_to_onehot = max_cat_to_onehot  # The threshold for deciding whether XGBoost should use one-hot encoding based split for categorical data
        self.max_cat_threshold = max_cat_threshold  # The maximum number of categories considered for each split
        self.eval_metric = eval_metric  # The metric used for monitoring the training result and early stopping
        self.early_stopping_rounds = early_stopping_rounds  # The early stopping activation
        self.callbacks = callbacks  # The list of callback functions that are applied at end of each iteration
        self.use_label_encoder = use_label_encoder  # The choice to use label encoder
        
    def evaluate_model(self, args, X, y, le_species, le_header):
        """
        Evaluate the model performance using cross-validation.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.

        Returns:
            None
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank information to the features
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Perform binarization on the features
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Perform normalization on the features
        std_acc = []
        avg_acc = 0.0  # Initialize variables to track the average accuracy
        n_samples = 0  # Initialize variables to track the number of samples
        split_assignments = get_split_assignments(args)  # Get assignments of samples to folds
        best_accuracy = -np.inf
        best_fold = -1
        start_evaluation = time.time()  # Start time of the evaluation
        for fold in range(args.n_folds):
            best_fold_accuracy = self.evaluate_fold(args, X, y, le_species, le_header, fold, split_assignments)  # Evaluate the model on the current fold
            if best_fold_accuracy > best_accuracy:
                best_accuracy = best_fold_accuracy  # Update the best accuracy
                best_fold = fold  # Update the best fold index
                print("\nFold {} did improve the model.".format(best_fold))
            else:
                print("\nFold {} did not improve the model.".format(fold))
            std_acc.append(best_fold_accuracy)  # Add the best fold accuracy to the list
            avg_acc += best_fold_accuracy  # Accumulate the accuracy
            n_samples += 1  # Accumulate the number of samples
        end_evaluation = time.time()  # End time of the evaluation
        time_evaluation = end_evaluation - start_evaluation  # Total evaluation time
        avg_acc /= n_samples  # Compute the average accuracy
        std_acc = torch.std(torch.FloatTensor(std_acc))  # Compute the standard deviation of accuracies
        print(f"\nMean {args.metric} and standard deviation after {args.n_folds} folds of {args.n_estimators} estimators: {avg_acc:.4f}% & {std_acc:.4f}.")
        print(f"Total time of the evaluation: {time_evaluation:.2f}s.")
        
    def train_model(self, args, X, y, le_species, le_header):
        """
        Train the model using the provided data.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.

        Returns:
            None
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank information to the features
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Perform binarization on the features
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Perform normalization on the features
        if args.endangered:
            eunis_red_list_crosswalks = get_eunis_red_list_crosswalks(args)
            endangered_red_list_habitats = get_endangered_red_list_habitats(args)
            endangered_eunis_habitats = get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header)
            X, y = add_endangered_habitats(X, y, endangered_eunis_habitats)  # Add endangered habitats to the data
        if args.augmentation:
            if not args.endangered:
                labels = None
            else:
                labels = endangered_eunis_habitats
            X, y = add_augmentation(X, y, labels)  # Perform data augmentation
        if args.dropout > 0:
            X = add_dropout(X, args.dropout, len(le_species.classes_))  # Apply dropout to the features
        if args.noise:
            X = add_noise(X, len(le_species.classes_))  # Add noise to the features
        if args.log:
            X = add_log(X, len(le_species.classes_))  # Apply logarithmic transformation to the features
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X = remove_features(args, X, features)  # Remove specified features from the input
        if args.standardization:
            X, scaler = add_standardization(X, scaler=None)  # Perform feature standardization
            set_scaler(args, scaler)  # Save the scaler for future use
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Get model parameters
        model = XGB(**model_parameters)  # Initialize the model
        print()
        start_training = time.time()  # Start time of the training
        model.fit(X, y, eval_set=[(X, y)])  # Train the model
        end_training = time.time()  # End time of the training
        time_training = end_training - start_training  # Total training time
        print(f"\nTotal time of the training: {time_training:.2f}s.")
        y_probas = model.predict_proba(X)  # Predict probabilities for the training data
        metric = get_metric(args, len(le_header.classes_))  # Get the evaluation metric
        accuracy = metric(torch.from_numpy(y_probas), torch.from_numpy(y)) * 100  # Compute the accuracy
        if model.best_iteration + args.early_stopping_rounds <= args.n_estimators:
            print(f'Early stopping occurred at iteration {model.best_iteration + args.early_stopping_rounds} with best iteration = {model.best_iteration} and best {args.metric} = {accuracy:.4f}%.')
        else:
            print(f'Early stopping dit not occur as best iteration = {model.best_iteration} and best {args.metric} = {accuracy:.4f}')
        set_xgb_model(args, model)  # Set the trained model for future use
        print("\nSuccessfully saved model at Models/XGB.json")
        
    def predict_model(self, args, X, y, le_species, le_header):
        """
        Perform predictions using the trained model.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.

        Returns:
            None
        """
        if args.rank:
            X = add_rank(X, len(le_species.classes_))  # Add rank information to the features
        if args.binarization:
            X = add_binarization(X, len(le_species.classes_))  # Perform binarization on the features
        if args.normalization:
            X = add_normalization(X, len(le_species.classes_))  # Perform normalization on the features
        if args.log:
            X = add_log(X, len(le_species.classes_))  # Apply logarithmic transformation to the features
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X = remove_features(args, X, features)  # Remove specified features from the input
        if args.standardization:
            scaler = get_scaler(args)  # Get the scaler for feature standardization
            X, _ = add_standardization(X, scaler=scaler)  # Perform feature standardization
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Get model parameters
        model = XGB(**model_parameters)  # Initialize the model
        model = get_xgb_model(args, model)  # Load the trained model
        start_prediction = time.time()  # Start time of the prediction
        predictions = model.predict(X)  # Perform predictions
        end_prediction = time.time()  # End time of the prediction
        time_prediction = end_prediction - start_prediction  # Total prediction time
        predictions = add_predictions_decoding(predictions, le_header)  # Decode the predicted labels
        set_predictions(args, predictions)  # Save the predictions
        print("\nSuccessfully saved predictions at Data/predictions.txt")
        print(f"Total time of the prediction: {time_prediction:.2f}s.")

    def evaluate_fold(self, args, X, y, le_species, le_header, fold, split_assignments):
        """
        Evaluate the model performance on a specific fold of the data.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
            fold (int): Fold number for evaluation.
            split_assignments (numpy.ndarray): Assigned data splits.

        Returns:
            float: Accuracy of the model on the given fold.
        """
        print('\n' + '*'*11) 
        print('* Fold: {} *'.format(fold))  # Print fold number
        print('*'*11 + '\n')
        X_train, X_test, y_train, y_test = add_fold_assignments(X, y, split_assignments, fold)  # Split data into train and test sets for the fold
        if args.endangered:
            eunis_red_list_crosswalks = get_eunis_red_list_crosswalks(args)  # Obtain EUNIS red list crosswalks
            endangered_red_list_habitats = get_endangered_red_list_habitats(args)  # Obtain endangered red list habitats
            endangered_eunis_habitats = get_endangered_eunis_habitats(eunis_red_list_crosswalks, endangered_red_list_habitats, le_header)  # Get endangered EUNIS habitats
            X_test, y_test = add_endangered_habitats(X_test, y_test, endangered_eunis_habitats)  # Add endangered habitats to the test data
        if args.augmentation:
            if not args.endangered:
                labels = None
            else:
                labels = endangered_eunis_habitats
            X_train, y_train = add_augmentation(X_train, y_train, labels)  # Perform data augmentation on the train data
        if args.dropout > 0:
            X_test = add_dropout(X_test, args.dropout, len(le_species.classes_))  # Apply dropout to the test data
        if args.noise:
            X_test = add_noise(X_test, len(le_species.classes_))  # Add noise to the test data
        if args.log:
            X_train = add_log(X_train, len(le_species.classes_))  # Apply logarithmic transformation to the train data
            X_test = add_log(X_test, len(le_species.classes_))  # Apply logarithmic transformation to the test data
        if args.features != ['all']:
            features = ''.join(args.features).replace(' ', '').split(',')
            X_train = remove_features(args, X_train, features)  # Remove specified features from the training set
            X_test = remove_features(args, X_test, features)  # Remove specified features from the test set
        if args.standardization:
            X_train, scaler = add_standardization(X_train, scaler=None)  # Perform feature standardization on the train data
            X_test, scaler = add_standardization(X_test, scaler=scaler)  # Perform feature standardization on the test data
        model_parameters = get_model_parameters(args, X_train.shape[1], len(le_header.classes_))  # Get model parameters
        model = XGB(**model_parameters)  # Initialize the model
        start_fold = time.time()  # Start time of the fold evaluation
        model.fit(X_train, y_train, eval_set=[(X_test, y_test)])  # Train the model on the train data and evaluate on the test data
        end_fold = time.time()  # End time of the fold evaluation
        time_fold = end_fold - start_fold  # Total fold evaluation time
        y_probas = model.predict_proba(X_test)  # Predict class probabilities for the test data
        metric = get_metric(args, len(le_header.classes_))  # Get the evaluation metric
        accuracy = metric(torch.from_numpy(y_probas), torch.from_numpy(y_test)) * 100  # Calculate accuracy using the evaluation metric
        if model.best_iteration + args.early_stopping_rounds <= args.n_estimators:
            print(f'\nEarly stopping occurred at iteration {model.best_iteration + 10} with best iteration = {model.best_iteration} and best {args.metric} = {accuracy:.4f}.')
        else:
            print(f'\nEarly stopping dit not occur as best iteration = {model.best_iteration} and best {args.metric} = {accuracy:.4f}')
        print(f"Time: {time_fold:.2f}s.")  # Print the time taken for the fold evaluation
        return accuracy

    def run(self, args, X, y, le_species, le_header):
        """
        Run the XGB model based on the specified pipeline.

        Args:
            args (argparse.Namespace): Arguments passed to the framework.
            X (scipy.sparse._csr.csr_matrix): Input features.
            y (numpy.ndarray): Target labels.
            le_species (sklearn.preprocessing._label.LabelEncoder): Label encoder for species.
            le_header (sklearn.preprocessing._label.LabelEncoder): Label encoder for habitat classes.
        """
        if args.pipeline == 'evaluation':
            self.evaluate_model(args, X, y, le_species, le_header)  # Run evaluation pipeline
        elif args.pipeline == 'training':
            self.train_model(args, X, y, le_species, le_header)  # Run training pipeline
        elif args.pipeline == 'prediction':
            self.predict_model(args, X, y, le_species, le_header)  # Run prediction pipeline

models/xgb.py
cesar-leblanc/hdm-framework
import pkg_resources
import requests
import os
import sys
import subprocess

REQUIRED_DEPENDENCIES = {  # Dictionary of required dependencies and their versions
    'captum': '0.6.0',
    'contextily': '1.3.0',
    'geopandas': '0.12.2', 
    'geopy': '2.3.0',
    'imbalanced-learn': '0.10.1',
    'joblib': '1.2.0',
    'matplotlib': '3.7.2',
    'notebook': '6.5.4',
    'numpy': '1.24.3',
    'pandas': '2.0.0',
    'plotly': '5.14.1',
    'pyproj': '3.5.0', 
    'pytopk': '0.1.0',
    'pytorch_tabnet': '4.1.0',
    'requests': '2.29.0',
    'rioxarray': '0.14.1', 
    'rtdl': '0.0.13',
    'rtree': '1.0.1',
    'scipy': '1.10.1',
    'shapely': '2.0.1', 
    'scikit-learn': '1.2.2',
    'torch':  '1.13.1',
    'torchmetrics': '0.11.4',
    'tqdm': '4.65.0',
    'xgboost': '1.7.5'
}

class Check:
    """
    Pipeline class for installation check.
    """
    def run(self, args):
        """
        Run the 'Check' pipeline of the framework.

        Args:
            args (argparse.Namespace): Parsed command-line arguments.
        """
        if args.check_files:  # Check if file checking is enabled
            github_api = 'https://api.github.com/repos/cesar-leblanc/hdm-framework/contents'  # GitHub API URL for repository contents
            stack = [(github_api, '')]  # Initialize stack for traversing repository contents
            missing_files = []  # List to store missing files
            while stack:  # Continue traversing while stack is not empty
                current_api, current_path = stack.pop()  # Get current API and path from stack
                try:
                    response = requests.get(current_api)  # Send GET request to GitHub API
                    if response.status_code != 200:  # Check if response was successful
                        print(f"\nFailed to retrieve GitHub repository contents. Error code: {response.status_code}.")
                    github_contents = response.json()  # Parse response as JSON
                    for file_info in github_contents:  # Iterate over repository contents
                        file_name = file_info["name"]  # Get file name
                        file_path = os.path.join(os.getcwd(), current_path, file_name)  # Get absolute file path
                        if file_info["type"] == "file":  # Check if it's a file
                            if not os.path.exists(file_path):  # Check if file doesn't exist locally
                                missing_files.append(file_path)  # Add missing file to list
                        elif file_info["type"] == "dir":  # Check if it's a directory
                            stack.append((file_info["url"], os.path.join(current_path, file_name)))  # Add directory to stack for further traversal
                except requests.exceptions.RequestException as e:
                    print("\nAn error occurred while accessing the GitHub API.")  # Handle API request error
                except OSError as e:
                    print("\nAn error occurred while accessing the local file system.")  # Handle local file system error
            if missing_files:  # Check if there are missing files
                print("\nMissing files:")
                for file_name in missing_files:  # Print each missing file
                    print(file_name)
            else:
                print("\nNo missing files.")  # Indicate that no files are missing
        else:
            print("\nNot checking if files are missing.")  # File checking is disabled

        if args.check_dependencies:  # Check if dependency checking is enabled
            missing_dependencies = {}  # Dictionary to store missing dependencies
            for dependency, required_version in REQUIRED_DEPENDENCIES.items():  # Iterate over required dependencies
                try:
                    pkg_resources.require(dependency)  # Check if the package is installed
                    installed_version = pkg_resources.get_distribution(dependency).version  # Get the installed version
                    if installed_version != required_version:  # Compare installed version with required version
                        missing_dependencies[dependency] = required_version  # Add missing dependency to the dictionary
                except (pkg_resources.DistributionNotFound, pkg_resources.VersionConflict):  # Handle package not found or version conflict
                    missing_dependencies[dependency] = required_version  # Add missing dependency to the dictionary
            if missing_dependencies:  # Check if there are missing dependencies
                print("\nMissing dependencies:")
                for dependency, version in missing_dependencies.items():  # Print each missing dependency and its required version
                    print(f"{dependency} - {version}")
            else:
                print("\nNo missing dependencies.")  # Indicate that no dependencies are missing
        else:
            print("\nNot checking if dependencies are missing.")  # Dependency checking is disabled
                
        if args.check_environment:  # Check if environment checking is enabled
            required_python_version = (3, 7)  # Define the minimum required Python version
            python_version = sys.version_info[:2]  # Get the current Python version
            if python_version < required_python_version:  # Compare Python version with required version
                print(f"\nPython {required_python_version[0]}.{required_python_version[1]} or later is required.")
            cuda_version = subprocess.getoutput("nvcc --version")  # Execute shell command to retrieve CUDA version
            if "command not found" in cuda_version:  # Check if CUDA is not installed
                print("\nCUDA is not installed.")
            else:
                if python_version >= required_python_version:  # Check if Python version meets the requirement
                    print("\nEnvironment is properly configured.")  # Indicate that the environment is properly configured
        else:
            print("\nNot checking if environment is properly configured.")  # Environment checking is disabled

pipelines/check.py
cesar-leblanc/hdm-framework
from data.load_data import get_eunis_habitats, get_eva_header, get_eva_species, get_altitudes, get_ecoregions, get_dunes, get_coasts, get_countries
from data.preprocess_data import add_header_encoding, add_gbif_normalization, add_species_encoding, add_spatial_split, add_input_data, add_target_values, add_longitude, add_latitude, add_altitude, add_country, add_ecoregion, add_dune, add_coast, add_header_decoding, add_species_decoding
from data.save_data import set_le_header, set_eva_to_gbif_species, set_le_species, set_split_assignment, set_ohe_country, set_ohe_ecoregion, set_ohe_dune, set_ohe_coast, set_eva_header, set_eva_species, set_input_data, set_target_values

class Dataset:
    """
    Pipeline class for dataset preparation.
    """
    def run(self, args):
        """
        Run the 'Dataset' pipeline of the framework.

        Args:
            args (argparse.Namespace): Parsed command-line arguments.
        """
        eunis_habitats = get_eunis_habitats(args)  # Retrieve the current list of EUNIS habitats
        eva_header = get_eva_header(args, eunis_habitats)  # Retrieve header data
        eva_header, le_header = add_header_encoding(eva_header)  # Encode target labels with value between 0 and n_classes-1
        set_le_header(args, le_header)  # Save the header encoder
        eva_species = get_eva_species(args, eva_header)  # Retrieve the species data
        if args.gbif_normalization:
            eva_species, eva_to_gbif_species = add_gbif_normalization(eva_species, None)  # Normalize species names using GBIF
            set_eva_to_gbif_species(args, eva_to_gbif_species)  # Save the mapping dictionary 
        eva_species, le_species = add_species_encoding(eva_species, None)  # Encode species labels
        set_le_species(args, le_species)  # Save the species encoder
        eva_header = eva_header[eva_header['PlotObservationID'].isin(eva_species['PlotObservationID'])]  # Filter header and species based on common PlotObservationID
        eva_header = eva_header.reset_index(drop=True)  # Reset the index
        eva_species = eva_species.reset_index(drop=True)  # Reset the index
        split_assignment, eva_header = add_spatial_split(args, eva_header)  # Perform spatial split assignment
        set_split_assignment(args, split_assignment)  # Save the split assignments
        X = add_input_data(eva_header, eva_species, le_species)  # Add input data (X)
        y = add_target_values(eva_header)  # Add target values (y)
        if args.location_features:
            X = add_longitude(X, eva_header)  # Add longitudes
            X = add_latitude(X, eva_header)  # Add latitudes
            altitudes = get_altitudes(args)  # Retrieve altitudes data
            X, eva_header = add_altitude(X, eva_header, altitudes)
            countries = get_countries(args)  # Retrieve countries data
            X, y, ohe_country, eva_header, eva_species = add_country(args, X, y, eva_header, eva_species, countries, None)  # Add countries
            set_ohe_country(args, ohe_country)  # Save the countries encoder
            ecoregions = get_ecoregions(args)  # Retrieve ecoregions data
            X, ohe_ecoregion, eva_header = add_ecoregion(X, eva_header, ecoregions, None)  # Add ecoregions
            set_ohe_ecoregion(args, ohe_ecoregion)  # Save the ecoregions encoder
            dunes = get_dunes(args)  # Retrieve dunes data
            X, ohe_dune, eva_header = add_dune(X, eva_header, dunes, None)  # Add dunes
            set_ohe_dune(args, ohe_dune)  # Save the dunes encoder
            coasts = get_coasts(args)  # Retrieve coasts data
            X, ohe_coast, eva_header = add_coast(X, eva_header, coasts, None)  # Add coasts
            set_ohe_coast(args, ohe_coast)   # Save the coasts encoder
        eva_species = add_species_decoding(eva_species, le_species)  # Decode species labels
        eva_header = add_header_decoding(eva_header, le_header)  # Decode header labels
        set_eva_header(args, eva_header)   # Save the header data
        set_eva_species(args, eva_species)  # Save the species data
        set_input_data(args, X)  # Save the input data
        set_target_values(args, y)  # Save the target values

pipelines/dataset.py
cesar-leblanc/hdm-framework
from models.mlp import MLP
from models.rfc import RFC
from models.tnc import TNC
from models.xgb import XGB
from models.ftt import FTT

from utils import get_model_parameters

from data.load_data import get_input_data, get_target_values, get_le_species, get_le_header

AVAILABLE_EVALUATIONS = {  # Dictionary of available models for evaluation
    "mlp": MLP, 
    "rfc": RFC,
    "tnc": TNC,
    "xgb": XGB,
    "ftt": FTT
}

class Evaluation:
    """
    Pipeline class for parameters evaluation.
    """
    def run(self, args):
        """
        Run the 'Evaluation' pipeline of the framework.

        Args:
            args (argparse.Namespace): Parsed command-line arguments.
        """
        X = get_input_data(args)  # Retrieve input data (X)
        y = get_target_values(args)  # Retrieve target values (y)
        le_species = get_le_species(args)  # Retrieve the label encoder for species
        le_header = get_le_header(args)  # Retrieve the label encoder for header

        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Get model parameters based on number of input features and number of output classes
        
        model = AVAILABLE_EVALUATIONS[args.model](**model_parameters)  # Create an instance of the selected model
        model.run(args, X, y, le_species, le_header)  # Run the model's evaluation

pipelines/evaluation.py
cesar-leblanc/hdm-framework
from models.mlp import MLP
from models.rfc import RFC
from models.tnc import TNC
from models.xgb import XGB
from models.ftt import FTT

from utils import get_model_parameters

from data.load_data import get_test_header, get_test_species, get_eva_to_gbif_species, get_le_species, get_altitudes, get_ohe_country, get_countries, get_ohe_ecoregion, get_ecoregions, get_ohe_dune, get_dunes, get_ohe_coast, get_coasts, get_le_header
from data.preprocess_data import add_gbif_normalization, add_species_encoding, add_input_data, add_longitude, add_latitude, add_altitude, add_country, add_ecoregion, add_dune, add_coast, add_species_decoding
from data.save_data import set_test_data, set_test_header, set_test_species

AVAILABLE_PREDICTIONS = {  # Dictionary of available models for predictions
    "mlp": MLP, 
    "rfc": RFC,
    "tnc": TNC,
    "xgb": XGB,
    "ftt": FTT
}

class Prediction:
    """
    Pipeline class for habitats prediction.
    """
    def run(self, args):
        """
        Run the 'Prediction' pipeline of the framework.

        Args:
            args (argparse.Namespace): Parsed command-line arguments.
        """
        test_header = get_test_header(args)  # Retrieve test header data
        test_species = get_test_species(args, test_header)  # Retrieve test species data
        if args.gbif_normalization:
            eva_to_gbif_species = get_eva_to_gbif_species(args)  # Retrieve GBIF species mapping
            test_species, _ = add_gbif_normalization(test_species, eva_to_gbif_species)  # Normalize test species names using GBIF
        le_species = get_le_species(args)  # Retrieve species encoder
        le_header = get_le_header(args)  # Retrieve header encoder
        test_species, _ = add_species_encoding(test_species, le_species)  # Encode test species labels
        test_header = test_header.reset_index(drop=True)  # Reset the index
        test_species = test_species.reset_index(drop=True)  # Reset the index
        X = add_input_data(test_header, test_species, le_species)  # Add input data (X)
        if args.location_features:
            X = add_longitude(X, test_header)  # Add longitudes
            X = add_latitude(X, test_header)  # Add latitudes
            altitudes = get_altitudes(args)  # Retrieve altitudes data
            X, test_header = add_altitude(X, test_header, altitudes)  # Add altitudes
            ohe_country = get_ohe_country(args)  # Retrieve the countries encoder
            countries = get_countries(args)  # Retrieve countries data
            X, _, _, test_header, test_species = add_country(args, X, None, test_header, test_species, countries, ohe_country)  # Add countries
            ohe_ecoregion = get_ohe_ecoregion(args)  # Retrieve the ecoregions encoder
            ecoregions = get_ecoregions(args)  # Retrieve ecoregions data
            X, _, test_header = add_ecoregion(X, test_header, ecoregions, ohe_ecoregion)  # Add ecoregions
            ohe_dune = get_ohe_dune(args)  # Retrieve the dunes encoder
            dunes = get_dunes(args)  # Retrieve dunes data
            X, _, test_header = add_dune(X, test_header, dunes, ohe_dune)  # Add dunes
            ohe_coast = get_ohe_coast(args)  # Retrieve the coasts encoder
            coasts = get_coasts(args)  # Retrieve coasts data
            X, _, test_header = add_coast(X, test_header, coasts, ohe_coast)  # Add coasts
        set_test_data(args, X)  # Save the test data
        test_species = add_species_decoding(test_species, le_species)  # Decode test species labels
        set_test_header(args, test_header)  # Save the test header data
        set_test_species(args, test_species)  # Save the test species data
        
        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Retrieve model parameters

        model = AVAILABLE_PREDICTIONS[args.model](**model_parameters)  # Initialize the prediction model
        model.run(args, X, None, le_species, le_header)  # Run the prediction model

pipelines/prediction.py
cesar-leblanc/hdm-framework
from models.mlp import MLP
from models.rfc import RFC
from models.tnc import TNC
from models.xgb import XGB
from models.ftt import FTT

from utils import get_model_parameters

from data.load_data import get_input_data, get_target_values, get_le_species, get_le_header

AVAILABLE_TRAININGS = {  # Dictionary of available models for training
    "mlp": MLP, 
    "rfc": RFC,
    "tnc": TNC,
    "xgb": XGB,
    "ftt": FTT
}

class Training:
    """
    Pipeline class for models training.
    """
    def run(self, args):
        """
        Run the 'Training' pipeline of the framework.

        Args:
            args (argparse.Namespace): Parsed command-line arguments.
        """
        X = get_input_data(args)  # Retrieve input data (X)
        y = get_target_values(args)  # Retrieve target values (y)
        le_species = get_le_species(args)  # Retrieve the label encoder for species
        le_header = get_le_header(args)  # Retrieve the label encoder for header

        model_parameters = get_model_parameters(args, X.shape[1], len(le_header.classes_))  # Get model parameters based on number of input features and number of output classes

        model = AVAILABLE_TRAININGS[args.model](**model_parameters)  # Create an instance of the selected model
        model.run(args, X, y, le_species, le_header)  # Run the model's training

pipelines/training.py
cesar-leblanc/hdm-framework
import torch
import random
import numpy as np
import pytopk
import torchmetrics
import geopy.distance
import sys
import os
import shutil
import scipy.sparse
import multiprocessing
import sklearn.utils
import pandas as pd
import warnings

def bool_type(s):
    """
    Converts a string to a boolean value.

    Args:
        s (str): The input string to convert.

    Returns:
        bool: The converted boolean value.

    Raises:
        ValueError: If the input string is not a valid boolean value.
    """
    if s.lower() in ('true', 't', 'yes', 'y', '1'):
        return True
    elif s.lower() in ('false', 'f', 'no', 'n', '0'):
        return False
    else:
        raise ValueError('Invalid boolean value: %s' % s)
    
def disable_warnings(args):
    """
    Disable warnings based on the provided argument.

    Args:
        args (argparse.Namespace): Namespace containing command-line arguments.

    Note:
        If the 'disable_warnings' argument is True, this function will ignore all warnings to avoid cluttering the output with warning messages during the execution of the code.
    """
    if args.disable_warnings:
        warnings.filterwarnings("ignore")

def disable_bytecode_generation(args):
    """
    Disables bytecode generation by removing existing bytecode cache directories.

    Args:
        args (argparse.Namespace): The arguments containing the bytecode generation flag.

    Note:
        Bytecode generation can be enabled or disabled based on the value of the write_bytecode flag in args.
    """
    if args.write_bytecode:  # Check if bytecode generation is enabled
        sys.dont_write_bytecode = False  # Allow bytecode generation
    else:  # Check if bytecode generation is disabled
        sys.dont_write_bytecode = True  # Disable bytecode generation
        for root, dirs, files in os.walk(".", topdown=False):  # Remove existing bytecode cache directories
            for dir in dirs:
                if dir == "__pycache__":  # Check if the directory is "__pycache__"
                    shutil.rmtree(os.path.join(root, dir))  # Remove the "__pycache__" directory and its contents

def print_parameters(args):
    """
    Prints the parameters if the print_parameters flag is set.

    Args:
        args (argparse.Namespace): The arguments containing the parameters.

    Prints:
        The parameters and their corresponding values.
    """
    if args.print_parameters:  # Check if the print_parameters flag is set
        parameters = vars(args)  # Get the dictionary of parameters
        print('Parameters:')  # Print the heading for the parameters
        for k, v in parameters.items():  # Iterate over the parameters
            print(k + ': ' + str(v))  # Print the parameter name and value

def set_seed(args):
    """
    Sets the random seeds for reproducibility.

    Args:
        args (argparse.Namespace): The arguments containing the seed values.

    Note:
        The seed is set differently based on the pipeline and model selected.
        If the pipeline is 'evaluation', 'training', or 'prediction' and the model is 'mlp' or 'ftt', the seed is set for the torch module.
    """
    random.seed(args.global_seed)  # Set the random seed for the random module
    np.random.seed(args.global_seed)  # Set the random seed for the numpy module
    if args.pipeline in ['evaluation', 'training', 'prediction'] and args.model in ['mlp', 'ftt']:
        torch.manual_seed(args.seed)  # Set the seed for the torch module
        if args.use_gpu:
            torch.cuda.manual_seed(args.seed)  # Set the seed for the torch CUDA module
    else:
        torch.manual_seed(args.global_seed)  # Set the seed for the torch module
        if args.use_gpu:
            torch.cuda.manual_seed(args.global_seed)  # Set the seed for the torch CUDA module

def get_device(args):
    """
    Retrieves the device for computation.

    Args:
        args (argparse.Namespace): The arguments containing the device configuration.

    Returns:
        torch.device: The device object representing the computation device.
    """
    if args.use_gpu:
        device = torch.device('cuda')  # Set the device to GPU
    else:
        device = torch.device('cpu')  # Set the device to CPU
    return device

def get_criterion(args, y_train):
    """
    Retrieves the criterion for loss calculation.

    Args:
        args (argparse.Namespace): The arguments containing the criterion configuration.
        y_train (numpy.ndarray): The training labels.

    Returns:
        torch.nn.modules.loss.CrossEntropyLoss | pytopk.noised_losses.BalNoisedTopK | pytopk.noised_losses.ImbalNoisedTopK: The criterion for loss calculation.
    """
    if args.model == 'mlp' or args.model =='ftt':
        if args.criterion == 'cross-entropy':
            criterion = torch.nn.CrossEntropyLoss()  # Use CrossEntropyLoss for criterion
        elif args.criterion == 'bal-noised-top-k':
            criterion = pytopk.BalNoisedTopK(k=args.k, epsilon=3.0)  # Use BalNoisedTopK for criterion
        elif args.criterion == 'imbal-noised-top-k': 
            criterion = pytopk.ImbalNoisedTopK(k=args.k, epsilon=0.01, max_m=0.7, cls_num_list=np.bincount(y_train).tolist(), scale=50, n_sample=5)  # Use ImbalNoisedTopK for criterion
    elif args.model == 'tnc':
        if args.loss_fn == 'cross-entropy':
            criterion = torch.nn.CrossEntropyLoss()  # Use CrossEntropyLoss for criterion
    return criterion

def get_optimizer(args, model):
    """
    Retrieves the optimizer for model parameter updates.

    Args:
        args (argparse.Namespace): The arguments containing the optimizer configuration.
        model (torch.nn.Module): The model to optimize.

    Returns:
        torch.optim.sgd.SGD | torch.optim.adam.Adam | torch.optim.adamw.AdamW: The optimizer for model parameter updates.
    """
    if args.optimizer == 'sgd':
        optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate_init, momentum=args.momentum, nesterov=args.nesterov, weight_decay=args.weight_decay)  # Use Stochastic Gradient Descent (SGD) optimizer with specified parameters
    elif args.optimizer == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate_init, weight_decay=args.weight_decay)  # Use Adam optimizer with specified parameters
    elif args.optimizer == 'adamw':
        optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate_init, weight_decay=args.weight_decay)  # Use AdamW optimizer with specified parameters
    return optimizer

def get_scheduler(args, optimizer):
    """
    Retrieves the scheduler for adjusting the learning rate.

    Args:
        args (argparse.Namespace): The arguments containing the scheduler configuration.
        optimizer (torch.optim): The optimizer for which the scheduler will adjust the learning rate.

    Returns:
        torch.optim.lr_scheduler.ReduceLROnPlateau: The scheduler for adjusting the learning rate.
    """
    if args.model == 'mlp' or args.model == 'ftt':
        if args.scheduler == 'plateau':
            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=args.factor, patience=args.patience)  # Use ReduceLROnPlateau scheduler with specified parameters
    elif args.model == 'tnc':
        if args.scheduler_fn == 'step':
            scheduler = torch.optim.lr_scheduler.StepLR
    return scheduler

def get_metric(args, num_classes):
    """
    Retrieves the metric for evaluating model performance.

    Args:
        args (argparse.Namespace): The arguments containing the metric configuration.
        num_classes (int): The number of classes for the metric.

    Returns:
        torchmetrics.classification.accuracy.Accuracy: The metric for evaluating model performance.
    """
    if args.metric == 'accuracy':
        metric = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes, average=args.average, top_k=args.k)  # Use Accuracy metric with specified parameters
    return metric

def get_available_subprocesses(args):
    """
    Get the number of available subprocesses (CPUs) based on the command-line arguments.

    Args:
        args (argparse.Namespace): Command-line arguments.

    Returns:
        int: The number of available subprocesses (CPUs).
    """
    if args.num_workers == -1:
        num_cpus = multiprocessing.cpu_count()  # Get the number of available CPUs
    else:
        num_cpus = args.num_workers  # Use the specified number of subprocesses from command-line arguments
    return num_cpus

class NormedLinear(torch.nn.Module):
    """
    Custom module for a normalized linear layer.
    """
    def __init__(self, in_features, out_features):
        """
        Initializes the NormedLinear module.

        Args:
            in_features (int): Number of input features.
            out_features (int): Number of output features.
        """
        super(NormedLinear, self).__init__()
        self.weight = torch.nn.Parameter(torch.Tensor(in_features, out_features))  # Define the weight matrix as a trainable parameter
        self.weight.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5)  # Initialize the weight matrix with random values and normalize it along dimension 1

    def forward(self, x):
        """
        Forward pass of the normalized linear layer.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor.
        """
        out = torch.nn.functional.normalize(x, dim=1).mm(torch.nn.functional.normalize(self.weight, dim=0))  # Normalize the input tensor along dimension 1 (row-wise) and perform matrix multiplication between the normalized input and normalized weight matrix
        return out

def coords_to_bin(x, y, x_bin_width, y_bin_width):
    """
    Convert coordinates to bin indices based on the specified bin widths.

    Args:
        x (numpy.ndarray): Array of x-coordinates.
        y (numpy.ndarray): Array of y-coordinates.
        x_bin_width (float): Width of the x-coordinate bins.
        y_bin_width (float): Width of the y-coordinate bins.

    Returns:
        tuple: Tuple containing two numpy arrays, (x_bin_list, y_bin_list), representing the bin indices.
    """
    assert np.all(x >= 0)  # Assert that all x-coordinates are non-negative
    assert np.all(y >= 0)  # Assert that all y-coordinates are non-negative
    assert x_bin_width > 0  # Assert that x_bin_width is positive
    assert y_bin_width > 0  # Assert that y_bin_width is positive
    x_bin_list = np.array(np.floor(x / x_bin_width), dtype=int)  # Compute x-coordinate bin indices
    y_bin_list = np.array(np.floor(y / y_bin_width), dtype=int)  # Compute y-coordinate bin indices
    return (x_bin_list, y_bin_list)

def lon_to_global_easting(lon, origin):
    """
    Convert longitude coordinates to global easting distances in kilometers.

    Args:
        lon (numpy.ndarray): Array of longitude coordinates.
        origin (tuple): Tuple containing the origin coordinates (latitude, longitude).

    Returns:
        numpy.ndarray: Array of global easting distances in kilometers.
    """
    WGS84_EQUATOR_CIRCUMFERENCE = 40075.017  # Equator circumference of WGS84 ellipsoid in kilometers
    global_eastings_km = np.zeros_like(lon)  # Initialize an array to store the global easting distances
    for i in range(len(lon)):
        lon_ofs = lon[i] - origin[1]  # Calculate the offset of longitude from the origin longitude
        if lon_ofs < -180:
            lon_ofs = 180 - (lon_ofs + 180)  # If the offset is less than -180 degrees, wrap it around to the range [-180, 180]
        elif lon_ofs > 180:
            lon_ofs = -180 - (lon_ofs - 180)  # If the offset is greater than 180 degrees, wrap it around to the range [-180, 180]
        global_eastings_km[i] = geopy.distance.geodesic(origin, (origin[0], lon[i])).km  # Calculate the geodesic distance between the origin and the current longitude coordinate, convert the distance to kilometers and store it in the global eastings array
        if lon_ofs < 0:
            global_eastings_km[i] = WGS84_EQUATOR_CIRCUMFERENCE - global_eastings_km[i]  # If the offset is negative, adjust the global easting distance by subtracting it from the equator circumference
    return global_eastings_km

def lat_to_global_northing(lat, origin):
    """
    Convert latitude coordinates to global northing distances.

    Args:
        lat (numpy.ndarray): Array of latitude coordinates.
        origin (tuple): Tuple containing the origin coordinates (latitude, longitude).

    Returns:
        numpy.ndarray: Array of global northing distances.
    """
    WGS84_MERIDIAN_CIRCUMFERENCE = 40007.863  # Circumference of the meridian in kilometers
    global_northings_km = np.zeros_like(lat)  # Initialize an array for global northing distances
    for i in range(len(lat)):
        lat_ofs = lat[i] - origin[0]  # Calculate the latitude offset from the origin
        if lat_ofs < -90:
            lat_ofs = 90 - (lat_ofs + 90)  # Adjust the latitude offset if it is less than -90 degrees
        elif lat_ofs > 90:
            lat_ofs = -90 - (lat_ofs - 90)  # Adjust the latitude offset if it is greater than 90 degrees
        global_northings_km[i] = geopy.distance.geodesic(origin, (lat[i], origin[1])).km  # Calculate the geodesic distance between the origin and the latitude coordinate, which represents the global northing distance in kilometers
        if lat_ofs < 0:
            global_northings_km[i] = WGS84_MERIDIAN_CIRCUMFERENCE - global_northings_km[i]  # Adjust the global northing distance if the latitude offset is negative
    return global_northings_km

def assign_block_ids(lon, lat, east_west_bin_km, north_south_bin_km, origin):
    """
    Assign block IDs to coordinates based on binning and origin.

    Args:
        lon (numpy.ndarray): Array of longitude coordinates.
        lat (numpy.ndarray): Array of latitude coordinates.
        east_west_bin_km (float): Width of the east-west bins in kilometers.
        north_south_bin_km (float): Height of the north-south bins in kilometers.
        origin (tuple): Tuple containing the origin coordinates (latitude, longitude).

    Returns:
        numpy.ndarray: Array of block IDs.
    """
    assert np.max(origin) <= 180  # Assert that the maximum value of origin is within [-180, 180]
    assert np.min(origin) >= -180  # Assert that the minimum value of origin is within [-180, 180]
    easting_list_km = lon_to_global_easting(lon, origin)  # Convert longitude coordinates to global easting distances
    northing_list_km = lat_to_global_northing(lat, origin)  # Convert latitude coordinates to global northing distances
    easting_bin_list, northing_bin_list = coords_to_bin(easting_list_km, northing_list_km, east_west_bin_km, north_south_bin_km)  # Convert the global easting and northing distances to bin indices
    block_id_list = []
    for i in range(len(easting_bin_list)):
        block_id_list.append("e" + str(easting_bin_list[i]).zfill(4) + "n" + str(northing_bin_list[i]).zfill(4))  # Generate block IDs by concatenating the easting and northing bin indices with appropriate formatting
    return np.array(block_id_list)

def get_model_parameters(args, n_features, n_outputs):
    """
    Get the model parameters based on the chosen model.

    Args:
        args (argparse.Namespace): Arguments object containing the framework configuration.
        n_features (int): Number of input features.
        n_outputs (int): Number of output classes.

    Returns:
        dict: Model parameters.
    """
    if args.model == 'mlp':
        model_parameters = {
            "input_dim": n_features,  # The number of input features
            "num_hidden_layers": args.num_hidden_layers,  # The number of hidden layers
            "num_neurons": args.num_neurons,  # The hidden dimension (number of neurons / number of non-linear activation functions)
            "output_dim": n_outputs,  # The number of output classes
            "activation_func": args.activation_func,  # The activtion function for the hidden layer
            "weight_init": args.weight_init,  # The choice of weight initializations
            "dropout_p": args.dropout_p,  # The dropout probability
            "use_batch_norm": args.use_batch_norm,  # The choice to enable or disable batch normalization
        }
    elif args.model == 'rfc':
        model_parameters = {
            "n_estimators": args.n_estimators,  # The number of trees in the forest
            "criterion": args.criterion,  # The function to measure the quality of a split
            "max_depth": args.max_depth,  # The maximum depth of the tree
            "min_samples_split": args.min_samples_split,  # The minimum number of samples required to split an internal node
            "min_samples_leaf": args.min_samples_leaf,  # The minimum number of samples required to be at a leaf node
            "min_weight_fraction_leaf": args.min_weight_fraction_leaf,  # The minimum weighted fraction of the sum total of weights (of all the input samples) required to be at a leaf node
            "max_features": args.max_features,  # The number of features to consider when looking for the best split
            "max_leaf_nodes": args.max_leaf_nodes,  # The value for which trees grow in best-first fashion
            "min_impurity_decrease": args.min_impurity_decrease,  # The value for which a node will be split if this split induces a decrease of the impurity greater than or equal to this value
            "bootstrap": args.bootstrap,  # The choice to use bootstrap samples when building trees
            "oob_score": args.oob_score,  # The choice to use out-of-bag samples to estimate the generalization score
            "n_jobs": args.n_jobs,  # The number of jobs to run in parallel
            "random_state": args.random_state,  # The value that controls both the randomness of the bootstrapping of the samples used when building trees and the sampling of the features to consider when looking for the best split at each node
            "verbose": args.n_estimators if args.verbose == 'n_estimators' else int(args.verbose),  # The value that controls the verbosity when fitting and predicting
            "warm_start": args.warm_start,  # The choice to reuse the solution of the previous call to fit and add mor estimators to the ensemble or to just fit a whole new forest
            "class_weight": args.class_weight,  # The weights associated with classes
            "ccp_alpha": args.ccp_alpha,  # The complexity parameter used for Minimal Cost-Complexity Pruning
            "max_samples": args.max_samples  # The number of samples to draw from X to train each base estimator if bootstrap is used
        }
    elif args.model == 'xgb':
        model_parameters = {
            "n_estimators": args.n_estimators,  # The number of boosting rounds
            "max_depth": args.max_depth,  # The maximum tree depth for base learners
            "max_leaves": args.max_leaves,  # The maximum number of leaves; 0 indicates no limit
            "max_bin": args.max_bin,  # The maximum number of bins per feature if using histogram-based algorithm
            "grow_policy": args.grow_policy,  # The tree growing policy
            "learning_rate": args.learning_rate,  # The boosting learning rate (xgb's 'eta')
            "verbosity": args.verbosity,  # The degree of verbosity
            "objective": args.objective,  # The learning task and the corresponding learning objective or a custom objective function to be used
            "booster": args.booster,  # The booster to use
            "tree_method": args.tree_method,  # The tree method to use
            "n_jobs": args.n_jobs,  # The number of parallel threads used to run xgboost
            "gamma": args.gamma,  # The minimum loss reduction required to make a further partition on a leaf node of the tree
            "min_child_weight": args.min_child_weight,  # The minimum sum of instance weight (hessian) needed in a child
            "max_delta_step": args.max_delta_step,  # The maximum delta step we allow each tree's weight estimation to be
            "subsample": args.subsample,  # The subsample ratio of the training instance
            "sampling_method": args.sampling_method,  # The sampling method
            "colsample_bytree": args.colsample_bytree,  # The subsample ratio of columns when constructing each tree
            "colsample_bylevel": args.colsample_bylevel,  # The subsample ratio of columns for each level
            "colsample_bynode": args.colsample_bynode,  # The subsample ratio of columns for each split
            "reg_alpha": args.reg_alpha,  # The L1 regularization term on weights (xgb's alpha)
            "reg_lambda": args.reg_lambda,  # The L2 regularization term on weights (xgb's lambda)
            "scale_pos_weight": args.scale_pos_weight,  # The balancing of positive and negative weights
            "base_score": args.base_score,  # The initial prediction score of all instances, global bias
            "random_state": args.random_state,  # The random number seed
            "missing": args.missing,  # The value in the data which needs to be present as a missing value
            "num_parallel_tree": args.num_parallel_tree,  # The value used for boosting random forest
            "monotone_constraints": args.monotone_constraints,  # The constraint of variable monotonicity
            "interaction_constraints": args.interaction_constraints,  # The constraints for interaction representing permitted interactions
            "importance_type": args.importance_type,  # The feature importance type for the feature importances property
            "gpu_id": args.gpu_id,  # The device ordinal
            "validate_parameters": args.validate_parameters,  # The choice to give warnings for unknown parameter
            "predictor": args.predictor,  # The value forces XGBoost to use specific predictor
            "enable_categorical": args.enable_categorical,  # The support for categorical data
            "feature_types": args.feature_types,  # The value used for specifying feature types without constructing a dataframe
            "max_cat_to_onehot": args.max_cat_to_onehot,  # he threshold for deciding whether XGBoost should use one-hot encoding based split for categorical data
            "max_cat_threshold": args.max_cat_threshold,  # The maximum number of categories considered for each split
            "eval_metric": args.eval_metric,  # The metric used for monitoring the training result and early stopping
            "early_stopping_rounds": args.early_stopping_rounds,  # The early stopping activation
            "callbacks": args.callbacks,  # The list of callback functions that are applied at end of each iteration
            "use_label_encoder": args.use_label_encoder  # The choice to use label encoder
        }
    elif args.model == 'tnc':
        model_parameters = {
            "n_d": args.n_d,  # The width of the decision prediction layer
            "n_a": args.n_a,  # The width of the attention embedding for each mask
            "n_steps": args.n_steps,  # The number of steps in the architecture (usually between 3 and 10)
            "gamma": args.gamma,  # The coefficient for feature reusage in the masks
            "cat_idxs": args.cat_idxs,  # The list of categorical features indices
            "cat_dims": args.cat_dims,  # The list of categorical features number of modalities (number of unique values for a categorical feature)
            "cat_emb_dim": args.cat_emb_dim,  # The list of embeddings size for each categorical features
            "n_independent": args.n_independent,  # The number of independent Gated Linear Units layers at each step
            "n_shared": args.n_shared,  # The number of shared Gated Linear Units at each step
            "epsilon": args.epsilon,  # The value should be left untouched
            "seed": args.seed,  # The random seed for reproducibility
            "momentum": args.momentum,  # The momentum for batch normalization
            "clip_value": args.clip_value,  # The gradient will be clipped at the given value
            "lambda_sparse": args.lambda_sparse,  # The extra sparsity loss coefficient
            "optimizer_fn": torch.optim.Adam if args.optimizer_fn == 'adam' else '?',  # The PyTorch optimizer function
            "optimizer_params": args.optimizer_params,  # The parameters compatible with the PyTorch optimizer function used to initialize the optimizer
            "scheduler_fn": get_scheduler(args, optimizer=None),  # The PyTorch scheduler to change learning rates during training
            "scheduler_params": args.scheduler_params,  # The dictionnary of parameters to apply to the PyTorch scheduler
            "model_name": args.model_name,  # The name of the model used for saving in disk
            "verbose": args.verbose,  # The verbosity for notebooks plots
            "device_name": args.device_name,  # The name of the device used for training
            "input_dim": n_features,  # The number of input features
            "output_dim": n_outputs,  # The number of output classes
            "mask_type": args.mask_type,  # The masking function to use for selection features
            "grouped_features": args.grouped_features,  # The groups to allow the model to share its attention across features inside a same group
            "n_shared_decoder": args.n_shared_decoder,  # The number of shared GLU block in decoder
            "n_indep_decoder": args.n_indep_decoder  # The number of independent GLU block in decoder
        }
    elif args.model == 'ftt':
        model_parameters = {
            "n_num_features": n_features,  # The number of input features
            "cat_cardinalities": args.cat_cardinalities,  # The number of unique values for each feature
            "d_token": args.d_token,  # The size of one token
            "n_blocks": args.n_blocks,  # The number of Transformer blocks
            "attention_n_heads": args.attention_n_heads,  # The number of attention heads
            "attention_dropout": args.attention_dropout,  # The dropout for attention blocks
            "attention_initialization": args.attention_initialization,  # The initialization for attention blocks
            "attention_normalization": args.attention_normalization,  # The normalization for attention blocks
            "ffn_d_hidden": args.ffn_d_hidden,  # The input size for the second linear layer in the Feed-Forward Network module
            "ffn_dropout": args.ffn_dropout,  # The dropout rate after the first linear layer in the Feed-Forward Network module
            "ffn_activation": args.ffn_activation,  # The activation used in the Feed-Forward Network
            "ffn_normalization": args.ffn_normalization,  # The normalization used in the Feed-Forward Network
            "residual_dropout": args.residual_dropout,  # The dropout rate rate for the output of each residual branch of all Transformer blocks
            "prenormalization": args.prenormalization,  # The choice to place normalizations at the beginning of each residual branch
            "first_prenormalization": args.first_prenormalization,  # The choice to keep the first normalization from the first Transformer layer
            "last_layer_query_idx": args.last_layer_query_idx,  # The indices of tokens that should be processed by the last Transformer block
            "n_tokens": None if int(args.kv_compression_ratio * n_features) == 0 else n_features + 1,  # The option for fast linear attention
            "kv_compression_ratio": None if int(args.kv_compression_ratio * n_features) == 0 else args.kv_compression_ratio,  # The choice to apply a technique to speed up attention modules when the number of features is large
            "kv_compression_sharing": None if int(args.kv_compression_ratio * n_features) == 0 else args.kv_compression_sharing,  # Weight sharing policy for the technique to speed up attention modules when the number of features is large
            "head_activation": args.head_activation,  # The activation used in the heads
            "head_normalization": args.head_normalization,  # The normalization used in the heads
            "d_out": n_outputs  # The number of output classes
        }
    return model_parameters

def get_fit_parameters(args, X_train, X_test, y_train, y_test):
    """
    Get the fit parameters based on the chosen model.

    Args:
        args (argparse.Namespace): Arguments object containing the framework configuration.
        X_train (scipy.sparse._csr.csr_matrix): Training input data.
        X_test (scipy.sparse._csr.csr_matrix): Testing input data.
        y_train (numpy.ndarray): Training target data.
        y_test (numpy.ndarray): Testing target data.

    Returns:
        dict: Fit parameters.
    """
    if args.model == 'tnc':
        fit_parameters = {
            'X_train': X_train,  # Training input data
            'y_train': y_train,  # Training target data
            'eval_set': [(X_test, y_test)],  # Evaluation set
            'eval_name': args.eval_name,  # The list of eval set names
            'eval_metric': args.eval_metric,  # The list of evaluation metrics
            'max_epochs': args.max_epochs,  # The maximum number of epochs for trainng
            'patience': args.patience,  # The number of consecutive epochs without improvement before performing early stopping
            'weights': args.weights,  # The sampling parameter
            'loss_fn': get_criterion(args, y_train),  # The loss function for training
            'batch_size': args.batch_size,  # The number of examples per batch
            'virtual_batch_size': args.virtual_batch_size,  # The size of the mini batches used for "Ghost Batch Normalization"
            'num_workers': get_available_subprocesses(args),  # The number or workers used in the data loader
            'drop_last': args.drop_last,  # The choice to drop last batch if not complete during training
            'callbacks': args.callbacks,  # The list of custom callbacks
            'pin_memory': args.pin_memory,  # Whether to set pin_memory to True or False during training
            'from_unsupervised': args.from_unsupervised,  # Use a previously self supervised model as starting weights
            'augmentations': args.augmentations,  # The choice to apply custom data augmentation pipeline during training
            'warm_start': args.warm_start,  # The choice to fit twice the same model and start from a warm start
            'compute_importance': args.compute_importance  # The choice to compute feature importance
        }
    return fit_parameters

class SparseDataset(torch.utils.data.Dataset):
    """
    Custom dataset class for sparse input data.
    """
    def __init__(self, X, y):
        """
        Initialize the SparseDataset.

        Args:
            X (scipy.sparse._csr.csr_matrix): Sparse input data.
            y (numpy.ndarray): Target data.
        """
        self.X = X  # Sparse input data
        self.y = y  # Target data

    def __getitem__(self, index):
        """
        Get a single item from the dataset.

        Args:
            index (int): Index of the item to retrieve.

        Returns:
            tuple: A tuple containing the input data and its corresponding target.
        """
        X = torch.from_numpy(self.X[index].toarray()[0]).float()  # Convert the sparse input data to a dense tensor
        y = self.y[index]  # Target data
        return X, y

    def __len__(self):
        """
        Get the total number of items in the dataset.

        Returns:
            int: Total number of items in the dataset.
        """
        return self.X.shape[0]  # Number of items in the dataset

utils.py
cesar-leblanc/hdm-framework