bmp_model.py

from dataclasses import dataclass, field
from typing import List, Tuple, Optional, Callable
import numpy as np
from scipy.optimize import minimize

# Define model A
def model_a(time, params):
    '''
    Simple first order BMP model
    y = B0 * (1 - exp(-k * t))'''
    B0, k = params
    return B0 * (1 - np.exp(-k * time))

# Define model B
def model_b(time, params):
    '''
    Two-phase BMP model
    y = B0 * (1 + (kH*exp(-kVFA*t)-kVFA*exp(-kH*t)/(kVFA-kH)))
    
    Note: after estimation of model A parameters:
        B0(0) = B0, kH = k, kVFA(0) = k*2
    '''  
    B0, kH, kVFA = params
    return B0 * (1 + (kH*np.exp(-kVFA*time)-kVFA*np.exp(-kH*time)/(kVFA-kH)))

# Define model C
def model_c(time, params):    
    '''
    Biphasic BMP model with fractional contribution
    y = B0 * (1 - alpha * exp(-kF * t) - (1 - alpha) * exp(-kL * t))

    Note: after estimation of model A parameters:
        B0(0) = B0, alpha(0) = 0.5, kF(0) = k, kL(0) = k/2
    '''  
    B0, alpha, kF, kL = params
    return B0 * (1 - alpha * np.exp(-kF * time) - (1 - alpha) * np.exp(-kL * time))

def model_d(time, params):
    '''
    Biphasic BMP model with fractional contribution and VFA inhibition
    y = B0 * (alpha*(1 + (kF*exp(-kVFA*t)-kVFA*exp(-kF*t))/(kVFA-kF)) + (1 - alpha) * (1 + (kL*exp(-kVFA*t)-kVFA*exp(-kL*t)/(kVFA-kL))))
    Note: after estimation of model A and model C parameters:
        B0(0) = B0, alpha(0) = 0.5, kF(0) = k, kL(0) = k/2, kVFA = kF*2 
    '''  
    B0, alpha, kF, kL, kVFA = params
    return B0 * (alpha*(1 + (kF*np.exp(-kVFA*time)-kVFA*np.exp(-kF*time))/(kVFA-kF)) + (1 - alpha) * (1 + (kL*np.exp(-kVFA*time)-kVFA*np.exp(-kL*time)/(kVFA-kL))))

def model_e(time, params):
    '''
    Modified Gompertz BMP model
    y = B0 * exp(-exp(1 + (lag - t)*(Rmax*exp(1)/B0)))
    ''' 
    B0, Rmax, lag = params
    return B0 * np.exp(-np.exp(1 + (lag - time)*(Rmax*np.exp(1)/B0)))

@dataclass
class BMP_estimation:
    """
    A class to evaluate and estimate parameters of a simple algebraic model.
    The model A is: y = param1 * (1 - exp(-param2 * t))
    To set the initial conditions of the models given an arbitrary first guess for model A, refer to: https://link.springer.com/article/10.1007/s00449-014-1150-4#appendices
    Model expressions are reported in the Appendix A of: https://www.sciencedirect.com/science/article/pii/S0960148122001562 
    """
    time: np.ndarray  # Time points
    data: Optional[np.ndarray] = None  # Observed data points (optional)
    model: Callable[[np.ndarray, List[float]], np.ndarray] = None  # Model function
    default_params: List[float] = field(default_factory=lambda: [1.0, 0.1])  # Default first guess for parameters [param1, param2]
    weights: Optional[np.ndarray] = None  # Weights for weighted least squares (optional)

    def evaluate_model(self, params: Optional[List[float]] = None) -> np.ndarray:
        """
        Evaluate the model with given parameters.

        Args:
            params (List[float]): Parameters for the model. If None, default parameters are used.

        Returns:
            np.ndarray: Model output for the given parameters.
        """
        if self.model is None:
            raise ValueError("No model function is defined.")
        if params is None:
            params = self.default_params
        return self.model(self.time, params)

    def compute_error(self, params: Optional[List[float]] = None) -> np.ndarray:
        """
        Compute the error between the model and observed data.

        Args:
            params (List[float]): Parameters for the model. If None, default parameters are used.

        Returns:
            np.ndarray: Error values for each time point.
        """
        if self.data is None:
            raise ValueError("Observed data is not provided.")
        model_output = self.evaluate_model(params)
        return self.data - model_output

    def compute_rmse(self, params: Optional[List[float]] = None) -> float:
        """
        Compute the Root Mean Squared Error (RMSE) between the model and observed data.

        Args:
            params (List[float]): Parameters for the model. If None, default parameters are used.

        Returns:
            float: RMSE value.
        """
        error = self.compute_error(params)
        rmse = np.sqrt(np.mean(error ** 2))
        return rmse

    def compute_weighted_rmse(self, params: Optional[List[float]] = None) -> float:
        """
        Compute the Weighted Root Mean Squared Error (WRMSE) between the model and observed data.

        Args:
            params (List[float]): Parameters for the model. If None, default parameters are used.

        Returns:
            float: Weighted RMSE value.
        """
        if self.weights is None:
            raise ValueError("Weights are not provided for weighted least squares.")
        error = self.compute_error(params)
        weighted_error = self.weights * (error ** 2)
        wrmse = np.sqrt(np.sum(weighted_error) / np.sum(self.weights))
        return wrmse

    def estimate_parameters(self, method: str = "BFGS", cost_function: Optional[Callable] = None) -> Tuple[List[float], float, Optional[np.ndarray]]:
        """
        Estimate the parameters by minimizing the specified cost function.

        Args:
            method (str): Optimization method to use (default: "BFGS").
            cost_function (Callable): Custom cost function to minimize. If None, defaults to RMSE or WRMSE.

        Returns:
            Tuple[List[float], float, Optional[np.ndarray]]:
                - Optimized parameters (rescaled to original magnitude).
                - Minimized cost function value.
                - Covariance matrix of the parameters (if available, otherwise None).
        """
        # Define scaling factors for normalization
        param_scales = np.array(self.default_params)
        scale_min = 1
        scale_max = 10

        def scale_params(params):
            """Scale parameters to the range (1, 10)."""
            return (params / param_scales) * (scale_max - scale_min) + scale_min

        def rescale_params(scaled_params):
            """Rescale parameters back to their original magnitude."""
            return (scaled_params - scale_min) / (scale_max - scale_min) * param_scales

        # Use RMSE or WRMSE as the default cost function
        if cost_function is None:
            if self.weights is None:
                cost_function = self.compute_rmse 
                print("Cost function and weights not provided, using RMSE")
            else:
                cost_function = self.compute_weighted_rmse
                print("Cost function not provided, using weights-RMSE")

        def wrapped_cost_function(scaled_params):
            """Cost function for optimization."""
            params = rescale_params(scaled_params)
            return cost_function(params)

        # Initial guess scaled
        scaled_initial_guess = scale_params(self.default_params)

        # Perform optimization
        result = minimize(wrapped_cost_function, scaled_initial_guess, method=method)

        # Rescale optimized parameters back to original magnitude
        optimized_params = rescale_params(result.x).tolist()
        minimized_cost = result.fun

        # Compute covariance matrix if available
        covariance_matrix = None
        if method == "BFGS" and result.hess_inv is not None:
            # Rescale Hessian inverse to original parameter space
            hess_inv = result.hess_inv
            scaling_factors = np.diag(param_scales / (scale_max - scale_min))
            covariance_matrix = scaling_factors @ hess_inv @ scaling_factors

        return optimized_params, minimized_cost, covariance_matrix