SteinFS: Feature Selection based on Stein's Formula

A Python package for nonparametric feature selection using Stein's formula, specifically designed for high-dimensional data with theoretical guarantees.

✨ Features

• Theoretical Guarantees: Rigorous mathematical foundation based on Stein's formula
• High-Dimensional: Efficiently handles high-dimensional data via screening
• Flexible: Supports both Gaussian and t-distributions
• Easy to Use: Clean scikit-learn-style API
• High Performance: Optimized NumPy and SciPy implementation

📦 Installation

From PyPI

pip install steinfs

🚀 Quick Start

Basic Usage

import numpy as np
from steinfs import SteinSelector

# Define covariance matrix
np.random.seed(42)
sigma = np.eye(200)  # Identity matrix (or use other covariance structures)

# Generate data from covariance matrix
X = np.random.multivariate_normal(mean=np.zeros(200), cov=sigma, size=2000)
y = X[:, 0]**2 + X[:, 1]**2 + X[:, 2]**2 + np.random.randn(2000) * 0.1  # True features are 0, 1, and 2

# Create selector and select top 5 features
selector = SteinSelector(num_features=5)
selector.fit(X, y, sigma=sigma)

# Get selected features
selected_features = selector.selected_features_
print(f"Selected features: {selected_features}")

# Transform data
X_selected = selector.transform(X)
print(f"Original shape: {X.shape}, Selected shape: {X_selected.shape}")

High-Dimensional Data with Screening

# Define covariance matrix
np.random.seed(42)
sigma = np.eye(2000)  # Identity matrix for high-dimensional data

# Generate high-dimensional data from covariance matrix
X_high = np.random.multivariate_normal(mean=np.zeros(2000), cov=sigma, size=2500)
y = X_high[:, 0]**2 + X_high[:, 5]**2 + X_high[:, 10]**2 + np.random.randn(2500) * 0.1 

# Use screening for feature selection
selector = SteinSelector(
    num_features=5,
    use_screening=True,  # Enable screening
    m=15,                # Number of screening iterations
    delta=0.9            # Keep 90% of features each time
)

X_selected = selector.fit_transform(X_high, y, sigma=sigma)
print(f"Selected features: {selector.selected_features_}")
print(f"Computation time: {selector.computation_time_:.4f}s")

t-Distribution Data

from scipy.stats import multivariate_t

# Define covariance matrix
np.random.seed(42)
nu = 5  # Degrees of freedom
sigma_t = np.eye(100)  # Covariance matrix for t-distribution

# Generate t-distribution data from covariance matrix
X_t = multivariate_t.rvs(loc=np.zeros(100), shape=sigma_t, df=nu, size= 2000)
y = X_t[:, 0]**2 + X_t[:, 1]**2 + X_t[:, 2]**2 + np.random.randn(2000) * 0.01 

# Use t-distribution selector
selector = SteinSelector(
    num_features=5,
    distribution='t',
    nu=5.0
)

X_selected = selector.fit_transform(X_t, y, sigma=sigma_t)
print(f"Selected features: {selector.selected_features_}")

Using Custom Covariance Matrix

# Define custom covariance matrix (e.g., identity matrix or Toeplitz structure)
np.random.seed(42)
sigma = np.eye(200)  # Identity matrix
# Or create Toeplitz structure: sigma[i,j] = rho^|i-j| for some rho

# Generate data from covariance matrix
X = np.random.multivariate_normal(mean=np.zeros(200), cov=sigma, size=2000)
y = X[:, 0]**2 + X[:, 1]**2 + X[:, 2]**2 + np.random.randn(2000) * 0.1

# Pass covariance matrix to fit method
selector = SteinSelector(num_features=5)
selector.fit(X, y, sigma=sigma)

print(f"Selected features: {selector.selected_features_}")
print(f"Covariance matrix shape: {selector.sigma_.shape}")

# Alternatively, use fit_transform with sigma
X_selected = selector.fit_transform(X, y, sigma=sigma)

🔧 API Reference

SteinSelector

Parameters:

• num_features (int, default=5): Number of features to select
• use_screening (bool, default=False): Whether to use screening
• m (int, default=10): Number of screening iterations
• delta (float, default=0.9): Proportion of features to keep in each screening step
• distribution (str, default='gaussian'): Distribution type ('gaussian' or 't')
• nu (float, default=5.0): Degrees of freedom for t-distribution
• rho (float, optional): Correlation coefficient for Toeplitz covariance
• random_state (int, optional): Random seed

Methods:

• fit(X, y, sigma=None): Fit the selector
• transform(X): Transform data, keeping only selected features
• fit_transform(X, y, sigma=None): Fit and transform
• get_support(indices=True): Get indices or mask of selected features

Attributes:

• selected_features_: Indices of selected features
• computation_time_: Computation time
• sigma_: Covariance matrix used

📚 Citation

If you use this package in your research, please cite our paper:

@article{du2025nonparametric,
  title={A Nonparametric Statistics Approach to Feature Selection in Deep Neural Networks with Theoretical Guarantees},
  author={Du, Junye and Li, Zhenghao and Gu, Zhutong and Feng, Long},
  year={2025},
  publisher={Oxford University Press}
}

🤝 Contributing

Contributions are welcome! Please feel free to submit issues or pull requests.

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

📧 Contact

• Junye Du: junyedu@connect.hku.hk

🙏 Acknowledgments

This work is supported by the Department of Statistics and Actuarial Science at The University of Hong Kong.

Keywords: Feature Selection, Stein's Formula, High-Dimensional Data, Nonparametric Statistics, Machine Learning