PSIDR.py

import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.manifold import MDS
from sklearn.utils import check_array, check_random_state
from sklearn.utils.validation import check_is_fitted
from scipy.optimize import minimize
from typing import Optional, Union
import numbers
from kbc import KBC



class PSIDR(BaseEstimator, TransformerMixin):
    """
    Point-distribution Similarity Invariant Dimensionality Reduction
    
    This estimator performs dimensionality reduction by:
    1. Clustering data
    2. Computing similarity-based distances between points and clusters
    3. Positioning cluster centers in low-dimensional space using MDS
    4. Determining point coordinates via multilateration
    
    Parameters
    ----------
    n_components : int, default=2
        Number of dimensions for the output embedding.
        
    k : int, default=10
        Number of clusters for KBC algorithm.
        
    tau : float, default=0.25
        Tau parameter for KBC (When using other clustering methods, please replace the parameters accordingly.).
        
    psi : int, default=64
        Psi parameter for KBC (When using other clustering methods, please replace the parameters accordingly.).
        
    random_state : int, RandomState instance or None, default=None
        Controls randomness in KBC and MDS.
        
    device : str, default='None'
        Device specification for KBC computation.
        
    Attributes
    ----------
    embedding_ : ndarray of shape (n_samples, n_components)
        Low-dimensional embedding of training data.
        
    labels_ : ndarray of shape (n_samples,)
        Cluster labels assigned to training data.
        
    class_centers_ : ndarray of shape (n_classes_, n_components)
        Low-dimensional coordinates of class centers.
        
    class_feature_map_ : ndarray of shape (n_classes_, psi)
        Mean feature map for each class.
        
    classes_ : ndarray of shape (n_classes_,)
        Unique class labels.
        
    n_classes_ : int
        Number of clusters found.
        
    class_offsets_ : ndarray of shape (n_classes_, n_components)
        Alignment offsets for each class center.
    """
    
    def __init__(self, n_components: int = 2, k: int = 10,
                 tau: float = 0.25, psi: int = 64,
                 random_state: Optional[Union[int, np.random.RandomState]] = None,
                 device: str = 'None'):
        # Initialize parameters without validation or computation
        self.n_components = n_components
        self.k = k
        self.tau = tau
        self.psi = psi
        self.random_state = random_state
        self.device = device
        
    def _validate_parameters(self):
        """Validate hyperparameters."""
        if not isinstance(self.n_components, numbers.Integral) or self.n_components < 1:
            raise ValueError(f"n_components must be a positive integer, got {self.n_components}")
        if not isinstance(self.k, numbers.Integral) or self.k < 1:
            raise ValueError(f"k must be a positive integer, got {self.k}")
        if not isinstance(self.tau, numbers.Real) or not (0 < self.tau < 1):
            raise ValueError(f"tau must be between 0 and 1, got {self.tau}")
        if not isinstance(self.psi, numbers.Integral) or self.psi < 1:
            raise ValueError(f"psi must be a positive integer, got {self.psi}")
            
    def fit(self, X, y=None):
        """
        Fit the model to training data.
        
        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
            
        y : Ignored
            Not used, present for API consistency.
            
        Returns
        -------
        self : object
            Returns the instance itself.
        """
        # Validate hyperparameters
        self._validate_parameters()
        
        # Validate input data
        X = check_array(X, dtype=np.float64, ensure_min_samples=2)
        
        # Set random state for reproducibility
        rng = check_random_state(self.random_state)
        
        # 1. Perform KBC clustering and extract feature map
        kbc = KBC(k=self.k, tau=self.tau, psi=self.psi,
                  random_state=rng, post_processing=True,
                  device=self.device)
        labels = kbc.fit_predict(X)
        feature_map = kbc.ndata  # Shape: (n_samples, psi)
        
        # 2. Compute class statistics
        classes, class_counts = np.unique(labels, return_counts=True)
        n_classes = len(classes)
        
        # Store class information
        self.classes_ = classes
        self.n_classes_ = n_classes
        
        # Calculate mean feature map for each class
        class_feature_map = np.zeros((n_classes, self.psi))
        for i, cls in enumerate(classes):
            mask = labels == cls
            class_feature_map[i] = feature_map[mask].mean(axis=0)
        
        # 3. Compute similarity matrices
        # Point-to-class similarity using inner product
        point_class_sim = feature_map @ class_feature_map.T  # (n_samples, n_classes)
        
        # Class-to-class similarity using inner product
        class_class_sim = class_feature_map @ class_feature_map.T  # (n_classes, n_classes)
        
        # 4. Convert similarities to distances and perform MDS on class centers
        # Distance formula: d = sqrt(2 * (1 - similarity))
        class_distances = np.sqrt(2 * (1 - class_class_sim))
        # Handle numerical instabilities
        class_distances = np.nan_to_num(class_distances, nan=1.0, posinf=1.0, neginf=1.0)
        
        mds = MDS(n_components=self.n_components,
                  dissimilarity='precomputed',
                  random_state=rng)
        self.class_centers_ = mds.fit_transform(class_distances)
        
        # 5. Determine point positions using multilateration
        point_class_distances = np.sqrt(2 * (1 - point_class_sim))
        point_class_distances = np.nan_to_num(point_class_distances, nan=1.0, posinf=1.0, neginf=1.0)
        
        Z = self._multilateration(point_class_distances, self.class_centers_)
        
        # 6. Align class centers with their point clouds
        self.class_offsets_ = np.zeros_like(self.class_centers_)
        for idx, cls in enumerate(classes):
            mask = labels == cls
            if np.any(mask):
                current_mean = Z[mask].mean(axis=0)
                self.class_offsets_[idx] = self.class_centers_[idx] - current_mean
                Z[mask] += self.class_offsets_[idx]
        
        # Store results as fitted attributes
        self.embedding_ = Z
        self.labels_ = labels
        self.class_feature_map_ = class_feature_map
        
        return self
    
    def _multilateration(self, distances, anchors):
        """
        Closed-form multilateration using least squares (MSE).
        
        Parameters
        ----------
        distances : ndarray of shape (n_points, n_anchors)
            Distances from each point to each anchor.
            
        anchors : ndarray of shape (n_anchors, n_dim)
            Coordinates of anchor points.
            
        Returns
        -------
        positions : ndarray of shape (n_points, n_dim)
            Estimated positions of points.
        """
        n_points, n_anchors = distances.shape
        n_dim = anchors.shape[1]
        
        # Use first anchor as reference for relative positioning
        A = -2 * (anchors[1:] - anchors[0:1])  # Shape: (n_anchors-1, n_dim)
        
        # Compute squared norms of anchors
        anchor_norms = np.sum(anchors ** 2, axis=1)
        
        # Setup least squares problem: A x = b
        # b_i = d_i^2 - d_0^2 - (||a_i||^2 - ||a_0||^2)
        b = (distances[:, 1:] ** 2 - distances[:, :1] ** 2) - (anchor_norms[1:] - anchor_norms[0])
        
        try:
            # Solve least squares problem
            # Using lstsq for robustness to singular matrices
            positions = np.linalg.lstsq(A, b.T, rcond=None)[0].T
            return positions
        except np.linalg.LinAlgError:
            # Fallback to optimization if least squares fails due to singular matrix
            positions = np.zeros((n_points, n_dim))
            x0 = anchors.mean(axis=0)
            
            for i in range(n_points):
                # Minimize sum of squared errors between distances
                res = minimize(
                    lambda x: np.sum((np.linalg.norm(anchors - x, axis=1) - distances[i]) ** 2),
                    x0,
                    method='BFGS',
                    options={'maxiter': 50}
                )
                positions[i] = res.x
                
            return positions
    
    def transform(self, X):
        """
        Transform new data to the low-dimensional embedding.
        
        Parameters
        ----------
        X : array-like of shape (n_samples_new, n_features)
            New data to transform.
            
        Returns
        -------
        X_new : ndarray of shape (n_samples_new, n_components)
            Low-dimensional embedding of new data.
        """
        check_is_fitted(self)
        
        # Validate input data
        X = check_array(X, dtype=np.float64)
        
        # Create new KBC instance for feature extraction from new data
        kbc = KBC(k=self.k, tau=self.tau, psi=self.psi,
                  random_state=self.random_state,
                  device=self.device)
        kbc.fit(X)
        feature_map_new = kbc.ndata
        
        # Compute similarities to training classes
        point_class_sim_new = feature_map_new @ self.class_feature_map_.T
        point_class_distances_new = np.sqrt(2 * (1 - point_class_sim_new))
        point_class_distances_new = np.nan_to_num(point_class_distances_new, nan=1.0, posinf=1.0, neginf=1.0)
        
        # Position new points using multilateration
        Z_new = self._multilateration(point_class_distances_new, self.class_centers_)
        
        # Align with predicted class centers
        pred_classes = np.argmax(point_class_sim_new, axis=1)
        for i in range(len(Z_new)):
            Z_new[i] += self.class_offsets_[pred_classes[i]]
            
        return Z_new
    
    def fit_transform(self, X, y=None):
        """
        Fit the model and return the transformed training data.
        
        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
            
        y : Ignored
            Not used, present for API consistency.
            
        Returns
        -------
        X_new : ndarray of shape (n_samples, n_components)
            Low-dimensional embedding of training data.
        """
        # Use fit().embedding_ for efficiency (avoid recomputing)

        return self.fit(X, y).embedding_