Source code for otlingam.models._greedy

from typing import ClassVar, Self

import numpy as np
from lingam.base import _BaseLiNGAM  # type: ignore
from sklearn.base import BaseEstimator  # type: ignore
from sklearn.utils._param_validation import validate_params  # type: ignore
from sklearn.utils.validation import validate_data  # type: ignore

from ..utils._wasserstein import gauss_quantiles




class GreedyOTLiNGAM(_BaseLiNGAM, BaseEstimator):
    """Greedy score-based causal discovery by sequential source removal.

    This estimator repeatedly selects the most non-Gaussian standardized residual as
    the next source in the causal order. It then removes the source's linear effect
    from every remaining variable. Once the ordering is recovered, edge weights are
    estimated using adaptive lasso regression.

    Data preprocessing settings:
        - `fit_intercept`: Whether to center the data before fitting. Centering also
          enables estimation of an intercept for each variable.

    Attributes:
        fit_intercept (bool): Whether to center the data before fitting.
        _causal_order (list[np.integer] | None): Internal causal ordering. None before
            fitting.
        _adjacency_matrix (np.ndarray | None): Internal weighted adjacency matrix.
            None before fitting.
        causal_order_ (list[np.integer]): Learned causal order from source to sink.
        adjacency_matrix_ (np.ndarray): Learned weighted adjacency matrix.
        intercept_ (np.ndarray): Intercepts of the regression models. Available only
            when `fit_intercept` is `True`.
        score_ (float): Sum of the selected squared Wasserstein scores.

    Examples:
        >>> from otlingam import GreedyOTLiNGAM
        >>> model = GreedyOTLiNGAM(fit_intercept=True)
        >>> model.fit(X)
        >>> model.causal_order_
    """

    fit_intercept: bool
    intercept_: np.ndarray
    score_: float

    _parameter_constraints: ClassVar[dict] = {"fit_intercept": ["boolean"]}



    def __init__(self, fit_intercept: bool = True):
        """Initializes GreedyOTLiNGAM.

        Args:
            fit_intercept (bool, optional): Whether to center the data. Defaults to
                True.
        """
        super().__init__()
        self.fit_intercept = fit_intercept




    @validate_params(
        {"X": ["array-like"], "y": [None]},
        prefer_skip_nested_validation=True,
    )
    def fit(self, X: np.typing.ArrayLike, y: None = None) -> Self:  # noqa: ARG002
        """Fits the GreedyOTLiNGAM algorithm.

        Args:
            X (np.typing.ArrayLike): Input data.
            y (None, optional): Ignored. Defaults to None.

        Returns:
            GreedyOTLiNGAM: The fitted estimator.

        Raises:
            ValueError: If a residual has zero variance.
        """
        self._validate_params()
        X = np.asarray(validate_data(self, X))  # type: ignore
        n, d = X.shape

        if self.fit_intercept:
            shift = X.mean(axis=0)
            X = X - shift

        residuals = X.copy()
        remaining = list(range(d))
        order = np.empty(d, dtype=int)
        quantiles = gauss_quantiles(n)
        score = 0.0

        for t in range(d):
            current = residuals[:, remaining]
            scales = np.sqrt(np.mean(current**2, axis=0))
            if np.any(scales == 0.0):
                raise ValueError("X must not contain a constant residual.")

            standardized = current / scales
            scores = np.mean(
                (np.sort(standardized, axis=0) - quantiles[:, None]) ** 2, axis=0
            )
            source_index = np.argmax(scores)
            source = remaining.pop(source_index)
            order[t] = source
            score += scores[source_index]

            if not remaining:
                break

            source_residual = residuals[:, source]  # type: ignore
            effects = (  # type: ignore
                source_residual
                @ residuals[:, remaining]
                / (source_residual @ source_residual)
            )
            residuals[:, remaining] -= np.outer(source_residual, effects)  # type: ignore

        self._causal_order = list(order)
        self._estimate_adjacency_matrix(X)
        self.score_ = float(score)

        if self.fit_intercept:
            self.intercept_ = shift - self._adjacency_matrix @ shift
        else:
            self.__dict__.pop("intercept_", None)

        return self