Source code for hetero.fitting

"""
This module defines the linear readout used to decode the target functions from the networks' state. This module is based on scikit-learn's linear model module and can be extended to other (possibly non-linear) regression methods as needed.
"""

import os 
osjoin = os.path.join
from pdb import set_trace

import numpy as np
from sklearn import linear_model
from sklearn.metrics import confusion_matrix, r2_score

from hetero.utils import timer 
    


class Decoder(object):
    """
    A decoder class for fitting and predicting using various regression algorithms.
    
    Parameters
    ----------
        optim_params : dict
            Dictionary containing optimization parameters:
            
            - 'alg': str, regression algorithm ('ridge', 'lasso', 'linreg', 'logreg')
            - 'fit_bias': bool, whether to include bias term
            - 'reg': float, regularization parameter (for ridge and lasso)
        
        ncpu : int, optional
            Number of CPU cores to use, by default -1 (all cores)
        
    
    .. warning:: The optimization algorithm must be compatible with the problem at hand. 
        No sanity checks are performed to assess that is the case or not.
    
    """
    
    def __init__(self, optim_params, ncpu=-1, **kwargs):
        
        self.alg = optim_params['alg']
        self.fit_bias = optim_params['fit_bias']
        self.alpha = optim_params['reg']
        self.n_jobs = ncpu
        self.score_fun = det_coeff
        self.coef_ = None
        
        # Standard ridge regression
        if self.alg == 'ridge':
            # if sag/saga solver is chosen we need to have
            # a fixed random state                      
            self.model = linear_model.Ridge(alpha=self.alpha, fit_intercept=self.fit_bias, 
                                            random_state=0 )
        
        # Standard lasso regression
        elif self.alg == 'lasso':
            self.model = linear_model.Lasso(alpha=self.alpha, fit_intercept=self.fit_bias)

            
        # Standard linear regression
        elif self.alg == 'linreg':
            self.model = linear_model.LinearRegression(fit_intercept=self.fit_bias, 
                                                       n_jobs=self.n_jobs)
            
        # Standard logistic regression
        elif self.alg == 'logreg':
            # if sag/saga solver is chosen we need to have
            # a fixed random state                      
            self.score_fun = conf_mat
            self.model = linear_model.LogisticRegression(fit_intercept=self.fit_bias, 
                                                         n_jobs=self.n_jobs, random_state=0)
        
        # Removed in favor of Scikit-learn that handles the intercept much nicer.
        ## direct least squares (scipy)
        # elif self.alg == 'direct':
        #     self.model = lstsq  
            
        
        else:
            raise ValueError(f'Fitting algorithm {self.alg} is invalid.')
       


    def fit(self, X, y):
        """Fit the model using features X and target y."""
        self.model.fit(X.T, y.T)

        


    def predict(self, X):
        """Predict using fitted model on features X."""
        return self.model.predict(X.T).T

    


    def score(self, yhat, y):
        """Compute score between predictions yhat and true values y."""
        return self.score_fun(yhat, y)

    
    def __get__(self):
        return self.model
        
    @property
    def coeff(self):
        """Get the model coefficients."""
        return self.model.coef_

    




def det_coeff(yhat, y, eps=1e-9):
    """Calculate the coefficient of determination (R-squared) score.

    Handles both single and multi-dimensional targets, with built-in outlier handling.
    For multi-dimensional data, computes R-squared separately for each dimension. 
    Relevant from regression problems.

    Parameters
    ----------
    yhat : array-like
        Predicted values
    y : array-like
        True values
    eps : float, optional
        Small constant to avoid division by zero, default 1e-9

    Returns
    -------
    float or array-like
        R-squared score(s). Returns array for multi-dimensional targets.
    """

    if y.ndim==2:
        v = np.sum((y - y.mean(axis=1, keepdims=True))**2, axis=1) + eps
        u = np.sum((y - yhat)**2, axis=1)
        return 1 - u/v
    else:
        return r2_score(y, yhat)





def conf_mat(yhat, y):
    """Compute confusion matrix between predicted and true labels. Relevant from classification
    problems.
    
    Parameters
    ----------
    yhat : array-like
        Predicted values
    y : array-like
        True values

    Returns
    -------
    2d array
        Confusion matrix
    """

    return confusion_matrix(y, yhat)