API Reference

def eval_all_indices(self):
    """Evaluates Sobol indices, Shapley effects, and total sensitivity indices.

    Uses the `indicies.eval_indices` utility.

    Updates the following attributes:
        self.results (pd.DataFrame): DataFrame containing Sobol indices results.
        self.non_zero_coefficients (pd.DataFrame): DataFrame of non-zero coefficients.
        self.shap (pd.DataFrame): DataFrame containing Shapley effects.
        self.total (pd.DataFrame): DataFrame containing total sensitivity indices.
    """
    eval_indicies = indicies.eval_indices(self.X_T_L, self.Y, self.coef_, self.evs) 
    self.results = eval_indicies.get_sobol_indicies()
    self.non_zero_coefficients = eval_indicies.get_non_zero_coefficients() 

    self.shap = eval_indicies.eval_shapley(self.X.columns)
    self.total = eval_indicies.eval_total_index(self.X.columns)

def get_deltax(self, num_unconditioned: int, delta_samples: int) -> pd.DataFrame:      
    """
    Calculate delta indices for the given number of unconditioned variables and delta samples.

    This method initializes a DeltaX instance using the provided data and parameters,
    then computes the delta indices based on the specified number of unconditioned variables
    and delta samples.

    Args:
        num_unconditioned (int): The number of unconditioned variables.
        delta_samples (int): The number of delta samples to generate.

    Returns:
        pd.DataFrame: A DataFrame containing the computed delta indices.
    """
    self.delta_instance = pawn.DeltaX(self.X, self.Y, self.ranges, self.non_zero_coefficients)
    delta_indices = self.delta_instance.get_deltax(num_unconditioned, delta_samples)
    return delta_indices

def get_hx(self, num_unconditioned: int, delta_samples: int) -> pd.DataFrame:      
    """
    Calculate delta indices for the given number of unconditioned variables and delta samples.

    This method initializes a DeltaX instance using the provided data and parameters,
    then computes the delta indices based on the specified number of unconditioned variables
    and delta samples.

    Args:
        num_unconditioned (int): The number of unconditioned variables.
        delta_samples (int): The number of delta samples to generate.

    Returns:
        pd.DataFrame: A DataFrame containing the computed delta indices.
    """
    self.delta_instance = pawn.hX(self.X, self.Y, self.ranges, self.non_zero_coefficients)
    delta_indices = self.delta_instance.get_hx(num_unconditioned, delta_samples)
    return delta_indices

def get_mc_shapley(self, joint=None, corr=None, N=10000,
                   method='exhaustive', n_perm=1000,
                   B=0, alpha=0.05, random_state=None,
                   f=None, progress=False, k_max=False):
    """Compute Shapley effects via Monte Carlo with correlated inputs.

    Uses a Monte Carlo approach to estimate Shapley effects when
    input variables may be correlated. The model function can be
    the trained surrogate model (default) or an arbitrary user-defined
    function.

    Two computation methods are available:
    - ``'exhaustive'``: Enumerates all 2^d - 1 non-empty subsets.
      Exact but computationally expensive for large d.
    - ``'permutation'``: Uses random permutations with lazy caching.
      More scalable for higher dimensions.

    Bootstrap confidence intervals are available for both methods.

    Args:
        joint: Distribution object with ``sample_joint(n)`` and
            ``sample_conditional(u_indices, fixed_x)`` methods.
            If ``None``, a :class:`GaussianCopulaUniform` is created
            from ``corr`` and the ranges of the training data.
        corr: Correlation matrix for the Gaussian copula, shape (d, d).
            Only used when ``joint is None``. Defaults to the identity
            matrix (independent inputs).
        N: Monte Carlo sample size per subset. Defaults to 10000.
        method: Computation method, ``'exhaustive'`` or
            ``'permutation'``. Defaults to ``'exhaustive'``.
        n_perm: Number of random permutations. Only used when
            ``method='permutation'``. Defaults to 1000.
        B: Number of bootstrap replications. Set to 0 to skip
            confidence intervals. Defaults to 0.
        alpha: Significance level for confidence intervals.
            Defaults to 0.05 (95% CI).
        random_state: Random seed for reproducibility.
        f: Model function ``f(x)`` where ``x`` is a 1D array and
            the return value is a scalar. If ``None``, the trained
            surrogate model's :meth:`predict` method is used.
            This allows both surrogate-based and user-defined
            function evaluations.
        progress: If ``True``, display tqdm progress bars during
            the Monte Carlo sampling (requires ``tqdm`` installed).
        k_max: Maximum coalition size for the exhaustive method.
            - ``False`` (default): auto-detect from surrogate's
              ``polys`` parameter (len(polys)).
            - ``None``: full enumeration (all 2^d − 1 subsets).
            - ``int``: explicit limit on coalition size.
            See the coalition truncation guide for details.

    Returns:
        pd.DataFrame: DataFrame with columns:
        - ``'variable'``: Variable names from the training data.
        - ``'effect'``: Normalised Shapley effects (sum to 1).
        - ``'shapley_value'``: Unscaled Shapley values.
        - ``'sobol_first'``: First-order Sobol indices S_i.
        - ``'sobol_total'``: Total-order Sobol indices T_i.
        - ``'total_variance'``: Estimated total output variance.
        - ``'lower'``, ``'upper'``: Bootstrap CI bounds (if B > 0).

    Examples:
        Using the trained surrogate model with independent inputs:

        >>> analyzer = rshdmr(data_file='data.csv', polys=[10, 5])
        >>> sobol, shapley, total = analyzer.run_all()
        >>> mc_results = analyzer.get_mc_shapley(N=5000, B=500)

        Using the surrogate model with a correlation matrix:

        >>> corr = np.array([[1.0, 0.5, 0.0],
        ...                  [0.5, 1.0, 0.0],
        ...                  [0.0, 0.0, 1.0]])
        >>> mc_results = analyzer.get_mc_shapley(corr=corr, N=5000)

        Using a user-defined function with a custom distribution:

        >>> def my_model(x):
        ...     return x[0]**2 + 2*x[1]
        >>> joint = MultivariateNormal(
        ...     mean=[0, 0], cov=[[1, 0.5], [0.5, 1]]
        ... )
        >>> mc_results = analyzer.get_mc_shapley(
        ...     joint=joint, f=my_model, N=5000
        ... )
    """
    # Determine the model function
    # self.predict handles both 1D (scalar) and 2D (batch) — pass as batch
    # predictor for efficient MC evaluation. User-defined f is wrapped for
    # compatibility and has no batch predictor.
    if f is None:
        _f = self.predict
        _predict_batch = self.predict
    else:
        _f = _wrap_predict_fn(f)
        _predict_batch = None

    # Build the distribution if not provided
    if joint is None:
        feature_names = list(self.X.columns)
        d = len(feature_names)
        lows = [self.ranges[name][0] for name in feature_names]
        highs = [self.ranges[name][1] for name in feature_names]
        if corr is None:
            corr = np.eye(d)
        joint = GaussianCopulaUniform(lows, highs, corr)

    # Auto-detect k_max from the surrogate model's polynomial orders.
    # k_max=False means "not set" → auto-detect.
    # k_max=None means "full enumeration".
    # k_max=int means explicit truncation.
    if k_max is False:
        if hasattr(self, 'polys') and self.polys:
            k_max = len(self.polys)
        else:
            k_max = None

    # Compute MC Shapley effects
    mc = MCShapley(_f, joint, predict_batch=_predict_batch)
    df = mc.compute(
        N=N, method=method, n_perm=n_perm,
        B=B, alpha=alpha, random_state=random_state,
        progress=progress, k_max=k_max,
    )

    # Replace generic variable names with actual feature names
    feature_names = list(self.X.columns)
    df['variable'] = feature_names

    return df

def get_pawn(self, S=10) :
    """Estimates PAWN sensitivity indices directly from data.

    Uses the `pawn.estimate_pawn` utility.

    Args:
        S (int, optional): Number of slices/intervals for the PAWN estimation. Defaults to 10.

    Returns:
        dict: Dictionary containing the PAWN sensitivity indices for each feature.
    """

    num_features = self.X.shape[1] 
    pawn_results = pawn.estimate_pawn(self.X.columns, num_features, self.X.values, self.Y, S=S)
    return pawn_results

def get_pawnx(self, num_unconditioned: int, num_conditioned: int, num_ks_samples: int, alpha: float = 0.05) -> pd.DataFrame:
    """Calculates PAWN sensitivity indices using the surrogate model.

    Uses the `pawn.pawnx` utility.

    Args:
        num_unconditioned (int): Number of unconditioned samples for PAWN.
        num_conditioned (int): Number of conditioned samples for PAWN.
        num_ks_samples (int): Number of samples for the Kolmogorov-Smirnov test.
        alpha (float, optional): Significance level for the KS test. Defaults to 0.05.

    Returns:
        pd.DataFrame: DataFrame containing the PAWN sensitivity indices.
    """
    pawn_instance = pawn.pawnx(self.X, self.Y, self.ranges, self.non_zero_coefficients)
    pawn_indices = pawn_instance.get_pawnx(num_unconditioned, num_conditioned, num_ks_samples, alpha) 
    return pawn_indices 

def get_pruned_data(self):
    """
    Generates a pruned dataset containing only the features with non-zero coefficients.

    This method creates a new DataFrame that includes only the columns from the original
    dataset (`X_T_L`) that correspond to the labels with non-zero coefficients. Additionally,
    it includes the target variable (`Y`).

    In streaming mode the active columns are computed on-the-fly from the lazy basis
    matrix instead of being sliced from a dense DataFrame.

    Returns:
        pd.DataFrame: A DataFrame containing the pruned data with selected features and the target variable.
    """
    pruned_data = pd.DataFrame()

    if self._streaming:
        # Build active columns from the lazy basis matrix
        active_labels = list(self.non_zero_coefficients['labels'])
        active_indices = [
            self._feature_names.index(label)
            for label in active_labels
        ]
        active_cols = self._lazy_basis.column_batch(active_indices)
        for i, label in enumerate(active_labels):
            pruned_data[label] = active_cols[:, i]
    else:
        for label in self.non_zero_coefficients['labels']:
            pruned_data[label] = self.X_T_L[label]

    pruned_data['Y'] = self.Y
    return pruned_data

def legendre_expand(self):
    """Performs Legendre polynomial expansion on the transformed data `self.X_T`.

    Uses the `legendre.legendre_expand` utility.

    For streaming methods ('omp_stream', 'omp_cv_stream'), builds a
    :class:`LazyBasisMatrix` instead of materialising the full design
    matrix, and creates a lightweight DataFrame with column labels only
    for downstream compatibility.

    Updates the following attributes:
        self.primitive_variables: Primitive variables from the expansion.
        self.poly_orders: Polynomial orders used in the expansion.
        self.X_T_L (pd.DataFrame): The expanded data matrix (dense) or a
            label-only DataFrame (streaming).
        self._lazy_basis: LazyBasisMatrix (streaming only).
    """
    self._streaming = self.method in ('omp_stream', 'omp_cv_stream')

    expansion_data = legendre.legendre_expand(self.X_T, self.polys)

    if self._streaming:
        self._lazy_basis = expansion_data.build_basis_set_streaming()
        # Create a lightweight DataFrame with column names only.
        # Downstream code (indicies.py, shapley.py) reads .columns
        # and never touches matrix values directly.
        self._feature_names = expansion_data._feature_names
        self.X_T_L = pd.DataFrame(columns=self._feature_names)
    else:
        expansion_data.build_basis_set()
        self.X_T_L = expansion_data.get_expanded()

    self.primitive_variables = expansion_data.get_primitive_variables()
    self.poly_orders = expansion_data.get_poly_orders()  

def predict(self, X):
    """Predicts output for new input data using the trained surrogate model.

    If the surrogate model (`self.surrogate_model`) doesn't exist, it first
    creates and fits one using `predictor.surrogate`.

    Args:
        X (array-like): Input data for which predictions are to be made.
            Should have the same features as the original training data.

    Returns:
        array-like: Predicted output values.
    """
    if not hasattr(self, 'surrogate_model'):
        self.surrogate_model = predictor.surrogate(self.non_zero_coefficients, self.ranges)
        self.surrogate_model.fit(self.X, self.Y)
    return self.surrogate_model.predict(X) 

def read_data(self, data_file):
    """Reads data from a file or DataFrame.

    Initializes the `self.X` (features) and `self.Y` (target) attributes.

    Args:
        data_file (str or pd.DataFrame): Path to the data file (CSV) or a pandas DataFrame.

    Raises:
        ValueError: If `data_file` is not a string path or a DataFrame.
    """
    if isinstance(data_file, pd.DataFrame):
        print('Found a DataFrame')
        df = data_file
    elif isinstance(data_file, str):
        df = pd.read_csv(data_file)
    else:
        raise ValueError("data_file must be either a pandas DataFrame or a file path (str).")

    self.Y = df['Y']
    self.X = df.drop('Y', axis=1)

    # Clean up the original DataFrame to save memory
    del df

def run_all(self):
    """
Execute a complete sequence of steps for RS-HDMR (Random Sampling High-Dimensional Model Representation) analysis.

This method performs the following steps in sequence:

1. Transforms the input data to a unit hypercube.
2. Builds basis functions using Legendre polynomials.
3. Runs regression analysis to fit the model.
4. Calculates and displays RS-HDMR model performance statistics.
5. Evaluates Sobol indices to quantify the contribution of each input variable to the output variance.
6. Calculates Shapley effects to measure the importance of each input variable.
7. Computes the total index to assess the overall impact of input variables.
8. If resampling is enabled, performs bootstrap resampling to estimate confidence intervals for Sobol indices and Shapley effects.
9. Prints a completion message with a randomly selected quote.

Returns:
    tuple: A tuple containing three elements:

        - sobol_indices (pd.DataFrame): A DataFrame containing Sobol indices for each input variable.
        - shapley_effects (pd.DataFrame): A DataFrame containing Shapley effects for each input variable.
        - total_index (pd.DataFrame): A DataFrame containing the total index for each input variable.

Notes:
    - The method assumes that the necessary data and configurations are already set in the class instance.
    - If resampling is enabled (`self.resampling` is True), confidence intervals (CIs) are calculated for Sobol indices and Shapley effects.
    - The method uses helper functions like `transform_data`, `legendre_expand`, `run_regression`, `stats`, `plot_hdmr`, `eval_sobol_indices`, `get_shapley`, and `get_total_index` to perform specific tasks.
    - The completion message includes a randomly selected quote for a touch of inspiration.

Example:
    >>> sobol_indices, shapley_effects, total_index = instance.run_all()
    >>> print(sobol_indices)
    >>> print(shapley_effects)
    >>> print(total_index)
"""
    # Define a helper function to print headings
    def print_step(step_name):
        print_heading(step_name)

    # Step 1: Transform data to unit hypercube
    print_step('Transforming data to unit hypercube')
    self.transform_data()

    # Step 2: Build basis functions
    print_step('Building basis functions')
    self.legendre_expand()

    # Step 3: Run regression analysis
    print_step('Running regression analysis')
    self.run_regression()

    # Step 4: Calculate RS-HDMR model performance statistics
    print_step('RS-HDMR model performance statistics')
    self.run_stats() 
    print()
    self.run_plot_hdmr()

    # Step 5: Evaluate Sobol indices
    self.eval_all_indices()
    sobol_indices = self.results.drop(columns=['labels', 'coeff'])

    # Step 6: Calculate Shapley effects
    shapley_effects = self.shap

    # Step 7: Calculate total index 
    total_index = self.total

    # Step 8: Perform resampling if enabled
    if self.resampling:
        print_step(f'Running bootstrap resampling {self.number_of_resamples} samples for {self.CI}% CI') 
        do_resampling = resample(self.get_pruned_data(), self.number_of_resamples, self.X.columns)
        do_resampling.do_resampling() 
        sobol_indices = do_resampling.get_sobol_quantiles(sobol_indices, self.CI)
        self.sobol_indices = sobol_indices.copy()
        # Calculate quantiles for Shapley effects
        shapley_effects = do_resampling.get_shap_quantiles(shapley_effects, self.CI) 
        print_step('Completed bootstrap resampling')

    # Step 9: Print completion message with a quote
    quote = quotes.get_quote() 
    message = (
        "                  Completed all analysis\n"
        "                 ------------------------\n\n"
        f"{textwrap.fill(quote, 58)}"
    )
    print_step(message)

    return sobol_indices, shapley_effects, total_index

def run_plot_hdmr(self):
    """
    Plots the High-Dimensional Model Representation (HDMR) of the model's predictions.

    This method uses the `plot_hdmr` function from the `stats` module to visualize the 
    HDMR of the actual values (`self.Y`) against the predicted values (`self.y_pred`).

    Returns:
        None
    """
    plot_hdmr(self.Y, self.y_pred)

def run_regression(self):
    """Runs the regression analysis using the specified method.

    Uses the `regression.regression` utility based on `self.method`.

    Updates the following attributes:
        self.coef_ (np.array): The regression coefficients obtained from the fit.
        self.y_pred (np.array): The predicted values based on the fitted model.
    """
    regression_instance = regression(
        X_T_L=self.X_T_L,
        Y=self.Y,
        method=self.method,
        n_iter=self.n_iter,
        verbose=self.verbose,
        cv_tol=self.cv_tol,
        starting_iter=self.starting_iter,
        cv_method=self.cv_method,
        n_jobs=self.n_jobs,
        ard_algorithm=self.ard_algorithm,
        threshold_lambda=self.threshold_lambda,
        cv_folds=self.cv_folds,
        return_std=self.return_std,
        ard_tol=self.ard_tol,
    )
    if self._streaming:
        self.coef_, self.y_pred = regression_instance.run_regression(
            lazy_basis=self._lazy_basis
        )
    else:
        self.coef_, self.y_pred = regression_instance.run_regression()

def run_stats(self):
    """Calculates and stores evaluation statistics for the fitted model.

    Uses the `stats.stats` utility.

    Updates the following attributes:
        self.evs (dict): A dictionary containing evaluation statistics (e.g., R^2, MSE).
    """
    self.evs = model_stats(self.Y, self.y_pred, self.coef_)

def transform_data(self):
    """Transforms the input data `self.X` into a unit hypercube.

    Updates the following attributes:
        self.ranges (list): The ranges (min, max) of the original data features.
        self.X_T (pd.DataFrame): The transformed data matrix within the unit hypercube.
    """
    transformed_data = transformation(self.X)
    transformed_data.do_transform()
    self.ranges = transformed_data.get_ranges()
    self.X_T = transformed_data.get_X_T()