MVPNA.py

#!/usr/bin/env python
# coding: utf-8

# In[ ]:

import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl
from torch.utils.data import TensorDataset, DataLoader
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay, roc_auc_score
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from sklearn.metrics import confusion_matrix, precision_score, recall_score
import seaborn as sns
#from tv import linear_operator_from_shape

def count_class(class_label, y_train):
    class_count =  len([x for x in y_train if x==class_label])
    return (2 * class_count) / len(y_train) 

#class_weights = torch.tensor([count_class(0) / count_class(1)], dtype=torch.float32)

#torch.manual_seed(0)
class FSGL(nn.Module):
    def __init__(
        self, input_size, learning_rate, adjacency_tensor,
        l1_lambda, l2_lambda, l1_lambda_inx, l2_lambda_inx, gl_lambda, fsl2_lambda, 
        groups_dict, ROI_cols, ROI_inx_cols, lower_triangular_cols, #tv_lambda=0.0, vertex_coords=None,
        use_inx=True, inx_terms_scaling_factor=1, num_classes=1, 
        loss_fn=nn.BCEWithLogitsLoss(pos_weight=torch.tensor([3.9344])), optimizer='Adam',
        patience=3,
        use_fm=False,
        fm_rank=4,
        fm_l2_lambda=0.0,
        fm_dropout=0.0,
        fm_importance_scale=1.0,
        fm_spatial_reg_weight=1.0,
        fm_init_std=None
    ):
    
        super(FSGL, self).__init__()
        
        self.groups_dict=groups_dict
        self.ROI_cols=ROI_cols
        self.ROI_inx_cols=ROI_inx_cols
        self.lower_triangular_cols=lower_triangular_cols
        self.optimizer = optimizer
        self.input_size = input_size
        
        feature_mask = torch.ones(input_size, dtype=torch.bool)
        # Mask subcortical vertices
        for i in [0, 35, 39]:
            feature_mask[self.groups_dict[i]] = False 
        if not use_inx:
            feature_mask[self.ROI_inx_cols] = False
        
        
        # Store cortical indices after mask (penalty data structures include subcortical vertices)
        self.cortical_feature_indices = torch.nonzero(feature_mask).squeeze().tolist()
        self.input_size_after_masking = len(self.cortical_feature_indices)

        # mapping from original indices to new indices after mask
        self.feature_index_map = {orig_idx: new_idx for new_idx, orig_idx in enumerate(self.cortical_feature_indices)}
        # linear layer to be Logisticreg
        self.linear = nn.Linear(self.input_size_after_masking, num_classes)
        
        ##### Factorization Machine inx layer
        # FM models all pairwise feature interactions in low-rank form without
        # explicitly constructing O(p^2) interaction columns.
        self.use_fm = use_fm
        self.fm_rank = int(fm_rank)
        self.fm_l2_lambda = float(fm_l2_lambda)
        self.fm_importance_scale = float(fm_importance_scale)
        self.fm_spatial_reg_weight = float(fm_spatial_reg_weight)

        if self.use_fm:
            if fm_init_std is None:
                # Important for very high-dimensional vertex data.
                # Prevents the FM term from exploding at initialization.
                fm_init_std = 1.0 / np.sqrt(
                    max(1, self.input_size_after_masking * self.fm_rank)
                )

            self.fm_V = nn.Parameter(
                torch.empty(self.input_size_after_masking, self.fm_rank)
            )
            nn.init.normal_(self.fm_V, mean=0.0, std=fm_init_std)

            self.fm_dropout = nn.Dropout(p=float(fm_dropout)) if fm_dropout > 0 else nn.Identity()
        else:
            self.register_parameter("fm_V", None)
            self.fm_dropout = nn.Identity()
        
        #nn.init.normal_(self.linear.weight, mean=1.0, std=0.01)
        self.output_layer = nn.Sigmoid()
        self.loss_fn = loss_fn
        self.learning_rate = learning_rate
        self.l1_lambda = l1_lambda
        self.l2_lambda = l2_lambda
        self.l1_lambda_inx = l1_lambda_inx
        self.l2_lambda_inx = l2_lambda_inx
        self.fsl2_lambda = fsl2_lambda
        self.gl_lambda = gl_lambda
        #self.tv_lambda = tv_lambda
        #self.vertex_coords = vertex_coords
        
        # tv linear_operator_from_shape() requires np array
        #if isinstance(vertex_coords, torch.Tensor):
        #    vertex_coords = vertex_coords.detach().cpu().numpy()

        # Subset vertex coords to match model's input after masking
        #cortical_coords = vertex_coords[self.cortical_feature_indices]
        #self.tv_operator = linear_operator_from_shape(cortical_coords)

        self.train_log = []
        self.val_log = []
        self.reg_log = []
        self.group_magnitudes = pd.DataFrame(columns = [i for i in range(12)])
        
        self.best_loss = np.inf
        self.patience = patience
        self.current_patience = 0
        self.early_stop = False
        self.use_inx = use_inx
        self.adjacency_tensor = adjacency_tensor
        self.inx_terms_scaling_factor = inx_terms_scaling_factor
        #self.class_weights = torch.tensor(class_weights, dtype=torch.float32) if class_weights is not None else None
        
        #self.mapped_adjacency = torch.tensor([[self.feature_index_map[idx1], self.feature_index_map[idx2]] 
        #                                 for idx1, idx2 in self.adjacency_tensor.tolist() if idx1 in self.feature_index_map and idx2 in self.feature_index_map])
        
        mapped_edges = [
            [self.feature_index_map[idx1], self.feature_index_map[idx2]]
            for idx1, idx2 in self.adjacency_tensor.tolist()
            if idx1 in self.feature_index_map and idx2 in self.feature_index_map
        ]

        if len(mapped_edges) == 0:
            self.mapped_adjacency = torch.empty((0, 2), dtype=torch.long)
            print("no mapped adj edges, ts chopped")
        else:
            self.mapped_adjacency = torch.tensor(mapped_edges, dtype=torch.long)
        
        #print('input_size arg:', self.input_size)
        #print('mask true count:', int(feature_mask.sum().item()))
        #print('self.input_size_after_masking:', self.input_size_after_masking)
        #print('expected masked-out:', int(self.input_size - self.input_size_after_masking))
        
    def forward(self, x):
        # Avoid in-place mutation of the batch tensor.
        if self.inx_terms_scaling_factor != 1:
            x = x.clone()
            x[:, self.lower_triangular_cols] *= self.inx_terms_scaling_factor

        x = x[:, self.cortical_feature_indices]

        logits = self.linear(x)

        if self.use_fm:
            logits = logits + self.fm_forward(x)

        # nn.BCEWithLogitsLoss() has built-in sigmoid
        if not isinstance(self.loss_fn, nn.BCEWithLogitsLoss):
            logits = self.output_layer(logits)

        return logits    
    '''
    def forward(self, x):
        
        #Scale inx terms
        if self.inx_terms_scaling_factor != 1:
            x[:, self.lower_triangular_cols] *= self.inx_terms_scaling_factor
        
        x = x[:, self.cortical_feature_indices]
        x = self.linear(x)
        # nn.BCEWithLogitsLoss() has built-in sigmoid
        if not isinstance(self.loss_fn, nn.BCEWithLogitsLoss):
            x = self.output_layer(x) 
        return x
    '''
    def fm_forward(self, x):
        """
        Second-order fm
        x: already-masked feature matrix, shape (n_samples, n_features_after_masking)

        Returns:
            shape (n_samples, 1), logit-scale interaction contribution.
        """
        if not self.use_fm:
            return torch.zeros((x.shape[0], 1), device=x.device, dtype=x.dtype)

        z = self.fm_dropout(x)
        xv = torch.matmul(z, self.fm_V) # (n, rank)
        x2v2 = torch.matmul(z * z, self.fm_V * self.fm_V)  # (n, rank)

        fm_term = 0.5 * torch.sum((xv * xv) - x2v2, dim=1, keepdim=True)
        return fm_term
    
    def fm_l2_penalty(self):
        if not self.use_fm:
            return torch.tensor(0.0, device=self.linear.weight.device)

        return self.fm_l2_lambda * torch.sum(self.fm_V ** 2)
    
    def regularization_loss(self):
        return (self.l1_penalty()
            + self.gl_penalty()
            + self.fsl_penalty()
            + self.elastic_net()
            + self.fm_l2_penalty())
    '''
    def l1_penalty(self):
        l1_norm = torch.tensor(0., dtype=torch.float32)
        l2_norm = torch.tensor(0., dtype=torch.float32)
        for weights in self.linear.weight:
            # Map the original vertex indices to new cortical vertex indices
            cortical_ROI_cols = []
            #subcortical_amt = 0
            #for idx in ROI_cols:
            #    if idx in self.feature_index_map:
            #        cortical_ROI_cols.append(self.feature_index_map[idx])         
            #    else:
            #        subcortical_amt+=1
            #print(subcortical_amt)
            ##Okay this works w cortical mapping
            cortical_ROI_cols = [self.feature_index_map[idx] for idx in self.ROI_cols if idx in self.feature_index_map]
            l1_norm += torch.linalg.vector_norm(weights[cortical_ROI_cols], ord=1)
            l2_norm += torch.linalg.vector_norm(weights[cortical_ROI_cols], ord=2)

        return self.l1_lambda * l1_norm + self.l2_lambda * l2_norm
    '''
    def l1_penalty(self):
        device = self.linear.weight.device

        l1_norm = torch.tensor(0.0, dtype=torch.float32, device=device)
        l2_norm = torch.tensor(0.0, dtype=torch.float32, device=device)

        cortical_ROI_cols = [
            self.feature_index_map[idx]
            for idx in self.ROI_cols
            if idx in self.feature_index_map
        ]

        for weights in self.linear.weight:
            l1_norm = l1_norm + torch.linalg.vector_norm(weights[cortical_ROI_cols], ord=1)
            l2_norm = l2_norm + torch.linalg.vector_norm(weights[cortical_ROI_cols], ord=2)

        return self.l1_lambda * l1_norm + self.l2_lambda * l2_norm
    
    def elastic_net(self):
        device = self.linear.weight.device

        if not self.use_inx:
            return torch.tensor(0.0, device=device)

        l1_norm = torch.tensor(0.0, dtype=torch.float32, device=device)
        l2_norm = torch.tensor(0.0, dtype=torch.float32, device=device)
        for weights in self.linear.weight:
            # Map original ROI_inx_cols bc subcortical vertices were removed
            mapped_ROI_inx_cols = [self.feature_index_map[idx] for idx in self.ROI_inx_cols]# if idx in self.feature_index_map]
            l1_norm += torch.linalg.vector_norm(weights[mapped_ROI_inx_cols], ord=1)
            l2_norm += torch.linalg.vector_norm(weights[mapped_ROI_inx_cols], ord=2)

        return self.l1_lambda_inx * l1_norm + self.l2_lambda_inx * l2_norm
        
    def print_param(self):
        count=0
        #for w in self.parameters():
        print('1: ',[w for w in self.parameters() if w.dim() > 1])
        print('\n2: ', print([torch.flatten(w, start_dim=1) for w in self.parameters() if w.dim() > 1]))
        print('3: ',[x for x in self.linear.weight])
        
    def gl_penalty(self, compute_group_magnitudes=False): #num_blk
            
        #store weight magnitudes of each group to plot over epochs
        #group_magnitudes_ls = []
        
        #gl_reg = torch.tensor(0., dtype=torch.float32)#.cuda() Can use with CUDA later
        device = self.linear.weight.device
        gl_reg = torch.tensor(0.0, dtype=torch.float32, device=device)
        
        for weights in self.linear.weight:
            if self.gl_lambda == 0.0:
                break
            
            for i in range(len(self.groups_dict.keys())):
                if i in [0, 39]:
                    continue
                
                # Map group indices to new cortical feature indices
                block_indices = [self.feature_index_map[idx] for idx in self.groups_dict[i] if idx in self.feature_index_map]
                #Tested, this works    
                block = weights[block_indices]  
                
                flat_block = torch.flatten(block, start_dim=0)
                gl_reg += np.sqrt(len(flat_block))*torch.norm(flat_block, p=2)#*(1.2/(torch.norm(flat_block, p=2)*10+0.0001))
                #group_magnitudes_ls.append(torch.sum(flat_block).detach().cpu().numpy())
        
        return gl_reg * self.gl_lambda
    
    def print_feature_weights(self, feature_names, return_dict=False, print_weights=True):
        
        if self.use_inx:
            if len(feature_names)+len(self.ROI_inx_cols)+2 != self.linear.weight.size(1):
                print(f'features: {len(feature_names)+len(self.ROI_inx_cols)+2}')
                print(f'weights: {self.linear.weight.size(1)}')
                raise ValueError("Num feature names != num weights.")
        else:
            if len(feature_names)+2 != self.linear.weight.size(1):
                print(f'features: {len(feature_names)+len(self.ROI_inx_cols)+2}')
                print(f'weights: {self.linear.weight.size(1)}')
                raise ValueError("Num feature names != num weights.")

        
        weights = self.linear.weight.detach().numpy().flatten()
        
        if return_dict:
            return dict(zip(feature_names, weights))
        
        if print_weights:
            print('Feature Weights:\n')
            for name, weight in zip(feature_names, weights):
                print(f"{name}: {weight}")
        return weights
        
    '''
    def fsl_penalty(self):
        
        if self.fsl2_lambda == 0.0:
            return torch.tensor(0., device=device)
        fused_norm = 0
        
        cortical_ROI_cols = [self.feature_index_map[idx] for idx in self.ROI_cols if idx in self.feature_index_map]
        ROI_dict = self.print_feature_weights(cortical_ROI_cols, return_dict=True)
        roi_array = torch.tensor([ROI_dict[col] for col in cortical_ROI_cols])
        
        for weights in self.linear.weight:
            region_diffs = weights[self.mapped_adjacency[:, 0]] - weights[self.mapped_adjacency[:, 1]]
            fused_norm += torch.sum(torch.abs(region_diffs))
        
        
        # Optional fused lasso on FM embeddings so interaction effects are spatially smooth.
        if self.use_fm and self.fm_spatial_reg_weight != 0.0:
            adj = self.mapped_adjacency.to(self.fm_V.device)
            fm_diffs = self.fm_V[adj[:, 0], :] - self.fm_V[adj[:, 1], :]
            fused_norm += self.fm_spatial_reg_weight * torch.sum(
                torch.linalg.vector_norm(fm_diffs, ord=2, dim=1)
            )        
                    
        return self.fsl2_lambda * fused_norm
    '''
    def fsl_penalty(self):
        device = self.linear.weight.device

        if self.fsl2_lambda == 0.0:
            return torch.tensor(0.0, device=device)

        fused_norm = torch.tensor(0.0, device=device)

        if self.mapped_adjacency.numel() == 0:
            return fused_norm

        adj = self.mapped_adjacency.to(device)

        # Linear/main-effect fused lasso
        for weights in self.linear.weight:
            region_diffs = weights[adj[:, 0]] - weights[adj[:, 1]]
            fused_norm = fused_norm + torch.sum(torch.abs(region_diffs))

        # FM fused lasso on latent interaction profiles
        if self.use_fm and self.fm_spatial_reg_weight != 0.0:
            fm_diffs = self.fm_V[adj[:, 0], :] - self.fm_V[adj[:, 1], :]
            fused_norm = fused_norm + self.fm_spatial_reg_weight * torch.sum(
                torch.linalg.vector_norm(fm_diffs, ord=2, dim=1)
            )

        return self.fsl2_lambda * fused_norm
    '''
    def tv_penalty(self):
        """
        Compute anisotropic Total Variation penalty using linear operator from ParsimonY.
        """
    
        if tv_lambda==0.0:
            return torch.tensor(0.)
        else:
            if vertex_coords==None:
                raise 'vertex_coords cannot be None if you want to compute tv_penalty'
        
        if self.tv_lambda == 0.0 or self.tv_operator is None:
            return torch.tensor(0.0, device=self.linear.weight.device)

        tv_pen = torch.tensor(0.0, dtype=torch.float32, device=self.linear.weight.device)

        for weights in self.linear.weight:
            # numpy is compatible with ParsimonY's linear operator
            w_np = weights.detach().cpu().numpy()

            # Apply linear operator (equivalent to gradient)
            grad_w = self.tv_operator.dot(w_np)

            # Anisotropic TV = L1 norm of gradients
            tv_pen += np.sum(np.abs(grad_w))

        return self.tv_lambda * tv_pen
    '''
        
    def fit(self, train_loader, val_loader, max_epochs, verbose=True):
        if self.optimizer == 'Adagrad':
            optimizer = torch.optim.Adagrad(self.parameters(), lr=self.learning_rate)
        else: 
            optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        #scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.2)
        #optimizer = torch.optim.LBFGS(self.parameters(), line_search_fn='strong_wolfe', lr=1)

        device = next(self.parameters()).device #move to GPU
        
        self.train()
        
        for epoch in range(max_epochs):
            
            train_loss = 0.0
            for batch in train_loader:
                x, y = batch
                x = x.to(device)
                y = y.to(device)
                
                optimizer.zero_grad()
                y_hat = self.forward(x)
                y = y.unsqueeze(1)  # Convert y to same shape as y_hat
                #y_hat = y_hat.unsqueeze(1)
                                
                loss = self.loss_fn(y_hat, y) + self.regularization_loss()
                self.train_log.append(self.loss_fn(y_hat, y).detach().cpu().numpy())
                if self.fsl2_lambda != 0.0:
                    self.reg_log.append(self.fsl_penalty().detach().cpu().numpy())
                elif self.gl_lambda != 0.0:
                    self.reg_log.append(self.gl_penalty().detach().cpu().numpy())
                elif self.l1_lambda != 0.0:
                    self.reg_log.append(self.l1_penalty().detach().cpu().numpy())
                
                loss.backward()
                optimizer.step()
                # PGD for lasso
                #with torch.no_grad():
                #    self.linear.weight.data = torch.sign(self.linear.weight.data) * torch.max(
                #        torch.zeros_like(self.linear.weight.data),
                #        torch.abs(self.linear.weight.data) - self.learning_rate*self.l1_lambda
                #    )
                # TODO: PGD for group and fused?
                
                
                train_loss += loss.item()
            
            # Every epoch check for early stopping
            valid_loss = 0.0
            if val_loader:
                with torch.no_grad():
                    self.eval()
                    for batch in val_loader:
                        x_val, y_val = batch
                        x_val = x_val.to(device)
                        y_val = y_val.to(device)
                        y_val_hat = self.forward(x_val)
                        y_val = y_val.unsqueeze(1)
                        #y_val_hat = y_val_hat.unsqueeze(1)
                        #print(y_val)
                        #print(y_val_hat)
                        loss_val = self.loss_fn(y_val_hat, y_val) + self.regularization_loss()
                        
                        self.val_log.append(self.loss_fn(y_val_hat, y_val).detach().cpu())
                        valid_loss += loss_val.item()
                        #valid_loss += self.loss_fn(y_val_hat, y_val).item()
                valid_loss /= len(val_loader)

            train_loss /= len(train_loader)
            
            if verbose:
                print(f'Epoch {epoch+1}/{max_epochs}, Train Loss: {train_loss:.4f}, Valid Loss: {valid_loss:.4f}')
            
            if val_loader:
                if round(valid_loss, 4) <= round(self.best_loss, 4):
                    self.best_loss = valid_loss
                    self.current_patience = 0
                    torch.save(self.state_dict(), 'best_model.pth')
                else:
                    self.current_patience += 1
                    if self.current_patience >= self.patience:
                        print(f'Early stopping at epoch {epoch+1}')
                        self.early_stop = True # For loading .pth of best weights if want
                        break

    # n plots
    """n_cols = len(self.group_magnitudes.columns)
        fig, axes = plt.subplots(4, int(n_cols/4), figsize=(20, 12))
        axes=axes.flatten()
        
        for i, col in enumerate(self.group_magnitudes.columns):
            axes[i].plot(self.group_magnitudes[col], label='weight magnitude')
            axes[i].axhline(y=0, color='black', linestyle='--')
            axes[i].legend()
            axes[i].set_title(f'Group {col}')
            axes[i].set_xlabel('Epoch')
            axes[i].set_ylabel('Weight vector magnitude')

        plt.tight_layout()
        plt.show()

        print(self.group_magnitudes)"""
    
    def get_full_weight_vector(self):
        full_weight = torch.zeros(self.input_size, device=self.linear.weight.device)
        for orig_idx, mapped_idx in self.feature_index_map.items():
            full_weight[orig_idx] = self.linear.weight[0][mapped_idx]
        return full_weight
    
    def get_full_fm_importance_vector(self):
        """
        Returns a full-length vector where each feature's value is the L2 norm
        of its FM embedding. This is an interaction-importance map.
        """
        full_importance = torch.zeros(self.input_size, device=self.linear.weight.device)

        if not self.use_fm:
            return full_importance

        fm_importance = torch.linalg.vector_norm(self.fm_V, ord=2, dim=1)

        for orig_idx, mapped_idx in self.feature_index_map.items():
            full_importance[orig_idx] = fm_importance[mapped_idx]

        return full_importance

    def get_full_importance_vector(
        self,
        include_linear=True,
        include_fm=True,
        normalize_components=True,
        eps=1e-8
    ):
        """
        Combined cluster-discovery map.

        By default:
            importance = z_+(|linear weight|) + fm_importance_scale * z_+(FM norm)

        z_+ means z-score over modeled cortical features, then clamp negatives to 0.
        This prevents the FM component from dominating only because of scale.
        """
        device = self.linear.weight.device

        full_mask = torch.zeros(self.input_size, dtype=torch.bool, device=device)
        idx = torch.tensor(self.cortical_feature_indices, dtype=torch.long, device=device)
        full_mask[idx] = True

        components = []

        if include_linear:
            linear_component = torch.abs(self.get_full_weight_vector())
            components.append(linear_component)

        if include_fm and self.use_fm:
            fm_component = self.fm_importance_scale * self.get_full_fm_importance_vector()
            components.append(fm_component)

        if len(components) == 0:
            return torch.zeros(self.input_size, device=device)

        if normalize_components:
            normed = []
            for comp in components:
                out = torch.zeros_like(comp)
                vals = comp[full_mask]
                vals_z = (vals - vals.mean()) / (vals.std(unbiased=False) + eps)
                out[full_mask] = torch.clamp(vals_z, min=0.0)
                normed.append(out)
            components = normed

        total = torch.zeros_like(components[0])
        for comp in components:
            total = total + comp

        return total
    
    def test(self, test_loader, return_fused_norm=True):
        
        #does it matter where this goes?
        device = next(self.parameters()).device
        
        self.eval()
        test_loss = 0
        correct = 0
        all_preds = []
        all_labels = []
        all_probs = []
        with torch.no_grad():
            for batch in test_loader:
                x, y = batch
                x = x.to(device)
                y = y.to(device)
                y_hat = self.forward(x)
                #if not isinstance(self.loss_fn, nn.BCEWithLogitsLoss):
                #    y_hat = torch.sigmoid(y_hat)
                #y_hat = y_hat.unsqueeze(1)
                y = y.unsqueeze(1)  # Convert y to same shape as y_hat
                loss = self.loss_fn(y_hat, y) + self.regularization_loss()
                test_loss += loss.item()
                if not isinstance(self.loss_fn, nn.BCEWithLogitsLoss):
                    y_pred = (y_hat > 0.5).float()
                else:
                    y_pred = (torch.sigmoid(y_hat) > 0.5).float()
                correct += (y_pred == y).float().sum().item()
                
                all_preds.extend(y_pred.cpu().numpy())
                all_labels.extend(y.cpu().numpy())
                all_probs.extend(y_hat.cpu().numpy()) #y_hat_probs ?
                
                if isinstance(self.loss_fn, nn.BCEWithLogitsLoss):
                    y_hat = torch.sigmoid(y_hat)
                
        test_loss /= len(test_loader.dataset)
        accuracy = correct / len(test_loader.dataset)
        precision = precision_score(all_labels, all_preds)
        recall = recall_score(all_labels, all_preds)
        
        cm = confusion_matrix(all_labels, all_preds)
        '''
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=['Negative', 'Positive'], yticklabels=['Negative', 'Positive'])
        plt.xlabel('Predicted')
        plt.ylabel('True')
        plt.show()
        '''
        auc = roc_auc_score(all_labels, all_probs)
        
        print(f'Test set: Average loss: {round(test_loss, 4)}, Accuracy: {round(accuracy, 4)}')
        print(f'Precision: {round(precision, 4)}, Recall: {round(recall, 4)}')
        print(f'AUC: {round(auc, 4)}') 
        
        if return_fused_norm:
            self.fsl2_lambda = 1.0
            fused_norm = self.fsl_penalty() / self.fsl2_lambda
            return accuracy, test_loss, y_hat, auc, fused_norm
        else:
            return accuracy, test_loss, y_hat, auc
        
def plot_convergence(train_loss, val_loss, reg_log):
    plt.figure(figsize=(5, 3))
    plt.plot(train_loss, label='train log loss')
    plt.plot(val_loss, label='val log loss')
    plt.plot(reg_log, label='regularization')
    plt.legend()
    plt.xlabel('Epoch')
    plt.ylabel('loss')
    plt.show()    



torch.manual_seed(0)
class LR_pytorch(nn.Module):
    def __init__(
        self,
        input_size,
        learning_rate,
        l1_lambda_inx,
        l2_lambda_inx,
        l1_lambda_linear,
        l2_lambda_linear,
        roi_inx_slice=None, # tuple(start, end) or None
        roi_linear_slice=None, # tuple(start, end) or None
        num_classes=1,
        loss_fn=nn.BCEWithLogitsLoss(pos_weight=torch.tensor([3.9344]))
    ):
        super(LR_pytorch, self).__init__()

        self.input_size = input_size
        self.linear = nn.Linear(self.input_size, num_classes)

        self.output_layer = nn.Sigmoid()
        self.loss_fn = loss_fn
        self.learning_rate = learning_rate

        self.l1_lambda_inx = l1_lambda_inx
        self.l2_lambda_inx = l2_lambda_inx
        self.l1_lambda_linear = l1_lambda_linear
        self.l2_lambda_linear = l2_lambda_linear

        self.roi_inx_slice = roi_inx_slice
        self.roi_linear_slice = roi_linear_slice

        self.train_log = []
        self.val_log = []
        self.reg_log = []
        self.group_magnitudes = pd.DataFrame(columns=[i for i in range(12)])
        
    def forward(self, x):
        x = self.linear(x)
        # nn.BCEWithLogitsLoss() has built-in sigmoid
        if not isinstance(self.loss_fn, nn.BCEWithLogitsLoss):
            x = self.output_layer(x) 
        return x
    
    def l1_penalty(self):
        reg = torch.tensor(0.0, device=self.linear.weight.device)

        # self.linear.weight shape: (1, n_features)
        w = self.linear.weight

        if self.roi_inx_slice is not None:
            s, e = self.roi_inx_slice
            w_inx = w[:, s:e]
            reg = reg + self.l1_lambda_inx * torch.linalg.vector_norm(w_inx, ord=1)
            reg = reg + self.l2_lambda_inx * torch.linalg.vector_norm(w_inx, ord=2)

        if self.roi_linear_slice is not None:
            s, e = self.roi_linear_slice
            w_linear = w[:, s:e]
            reg = reg + self.l1_lambda_linear * torch.linalg.vector_norm(w_linear, ord=1)
            reg = reg + self.l2_lambda_linear * torch.linalg.vector_norm(w_linear, ord=2)

        return reg
    
    def print_param(self):
        count=0
        #for w in self.parameters():
        print('1: ',[w for w in self.parameters() if w.dim() > 1])
        print('\n2: ', print([torch.flatten(w, start_dim=1) for w in self.parameters() if w.dim() > 1]))
        print('3: ',[x for x in self.linear.weight])
    
    def print_feature_weights(self, feature_names, return_dict=False, print_weights=True):
        
        if self.use_inx:
            if len(feature_names)+len(ROI_inx_cols)+2 != self.linear.weight.size(1):
                raise ValueError("Num feature names != num weights.")
        else:
            if len(feature_names)+2 != self.linear.weight.size(1):
                raise ValueError("Num feature names != num weights.")

        
        weights = self.linear.weight.detach().numpy().flatten()
        
        if return_dict:
            return dict(zip(feature_names, weights))
        
        if print_weights:
            print('Feature Weights:\n')
            for name, weight in zip(feature_names, weights):
                print(f"{name}: {weight}")
        return weights
    
    def fit(self, x, y, x_val, y_val, max_epochs, verbose=True):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        #scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.2)
        #optimizer = torch.optim.LBFGS(self.parameters(), line_search_fn='strong_wolfe', lr=1)

        self.train()
        
        for epoch in range(max_epochs):
            
            train_loss = 0.0

            #x, y = 

            optimizer.zero_grad()
            y_hat = self.forward(x)
            #y = y.view(-1, 1)      # force y to match y_hat shape
            #print(y.shape)
            #y = y.unsqueeze(1)  # Convert y to same shape as y_hat
            #print(y.shape)
            #y_hat = y_hat.unsqueeze(1)

            loss = self.loss_fn(y_hat, y) + self.l1_penalty()
            self.train_log.append(self.loss_fn(y_hat, y).detach().numpy())
            
            loss.backward()
            optimizer.step()
            # PGD for lasso
            #with torch.no_grad():
            #    self.linear.weight.data = torch.sign(self.linear.weight.data) * torch.max(
            #        torch.zeros_like(self.linear.weight.data),
            #        torch.abs(self.linear.weight.data) - self.learning_rate*self.l1_lambda
            #    )
            # TODO: PGD for group and fused?


            train_loss += loss.item()

            # Every epoch check for early stopping
            valid_loss = 0.0
            '''with torch.no_grad():
                self.eval()

                #x_val, y_val = 
                y_val_hat = self.forward(x_val)
                #y_val = y_val.unsqueeze(1)
                #y_val_hat = y_val_hat.unsqueeze(1)
                #print(y_val)
                #print(y_val_hat)
                loss_val = self.loss_fn(y_val_hat, y_val) + self.l1_penalty()
                self.val_log.append(self.loss_fn(y_val_hat, y_val))
                valid_loss += loss_val.item()
                #valid_loss += self.loss_fn(y_val_hat, y_val).item()
'''
            train_loss /= len(x)
            #valid_loss /= len(x_val)
            
            if verbose:
                print(f'Epoch {epoch+1}/{max_epochs}, Train Loss: {train_loss:.4f}, Valid Loss: {valid_loss:.4f}')

    def test(self, x, y):
        self.eval()
        test_loss = 0
        correct = 0
        all_preds = []
        all_labels = []
        all_probs = []
        with torch.no_grad():
            y_hat = self.forward(x)
            #if not isinstance(self.loss_fn, nn.BCEWithLogitsLoss):
            #    y_hat = torch.sigmoid(y_hat)
            #y_hat = y_hat.unsqueeze(1)
            #y = y.unsqueeze(1)  # Convert y to same shape as y_hat
            loss = self.loss_fn(y_hat, y) + self.l1_penalty()
            test_loss += loss.item()
            if not isinstance(self.loss_fn, nn.BCEWithLogitsLoss):
                y_pred = (y_hat > 0.5).float()
            else:
                y_pred = (torch.sigmoid(y_hat) > 0.5).float()
            correct += (y_pred == y).float().sum().item()

            all_preds.extend(y_pred.cpu().numpy())
            all_labels.extend(y.cpu().numpy())
            all_probs.extend(y_hat.cpu().numpy()) #y_hat_probs ?

            if isinstance(self.loss_fn, nn.BCEWithLogitsLoss):
                y_hat = torch.sigmoid(y_hat)
                
        test_loss /= len(x)
        accuracy = correct / len(x)
        precision = precision_score(all_labels, all_preds)
        recall = recall_score(all_labels, all_preds)
        
        #cm = confusion_matrix(all_labels, all_preds)
        
        #sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=['Negative', 'Positive'], yticklabels=['Negative', 'Positive'])
        #plt.xlabel('Predicted')
        #plt.ylabel('True')
        #plt.show()
        
        auc = roc_auc_score(all_labels, all_probs)
        
        print(f'Test set: Average loss: {round(test_loss, 4)}, Accuracy: {round(accuracy, 4)}')
        print(f'Precision: {round(precision, 4)}, Recall: {round(recall, 4)}')
        print(f'AUC: {round(auc, 4)}') 
        return accuracy, test_loss, y_hat, auc
        
    def get_full_weight_vector(self):
        full_weight = torch.zeros(self.input_size, device=self.linear.weight.device)
        for orig_idx, mapped_idx in self.feature_index_map.items():
            full_weight[orig_idx] = self.linear.weight[0][mapped_idx]
        return full_weight