l0-python 0.1.0__tar.gz → 0.1.1__tar.gz

{l0_python-0.1.0 → l0_python-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: l0-python
-Version: 0.1.0
+Version: 0.1.1
 Summary: L0 regularization for sparse neural networks and intelligent sampling
 Author-email: PolicyEngine <hello@policyengine.org>
 License: MIT License
@@ -41,6 +41,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0.0
 Requires-Dist: numpy>=1.20.0
+Requires-Dist: scipy>=1.7.0
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov>=3.0; extra == "dev"
@@ -49,6 +50,7 @@ Requires-Dist: mypy>=0.900; extra == "dev"
 Requires-Dist: ruff>=0.0.280; extra == "dev"
 Requires-Dist: pre-commit>=3.0; extra == "dev"
 Requires-Dist: build>=1.3; extra == "dev"
+Requires-Dist: scikit-learn>=1.0; extra == "dev"
 Provides-Extra: docs
 Requires-Dist: mystmd>=1.3; extra == "docs"
 Requires-Dist: jupyter>=1.0; extra == "docs"

{l0_python-0.1.0 → l0_python-0.1.1}/l0/__init__.py

@@ -26,6 +26,7 @@ from .penalties import (
     get_sparsity_stats,
     update_temperatures,
 )
+from .sparse import SparseL0Linear
 
 __all__ = [
     # Distributions
@@ -36,6 +37,8 @@ __all__ = [
     "L0DepthwiseConv2d",
     "SparseMLP",
     "prune_model",
+    # Sparse
+    "SparseL0Linear",
     # Gates
     "L0Gate",
     "SampleGate",

l0_python-0.1.1/l0/sparse.py

@@ -0,0 +1,321 @@
+"""
+L0 regularized linear models for sparse data.
+
+Implements L0-regularized linear regression that efficiently handles
+sparse input matrices without converting them to dense format.
+"""
+
+# ruff: noqa: N803 N806
+
+
+import numpy as np
+import torch
+import torch.nn as nn
+from scipy import sparse as sp
+
+
+class SparseL0Linear(nn.Module):
+    """
+    L0-regularized linear model for sparse input data.
+
+    Efficiently handles sparse matrices (scipy.sparse) without
+    converting to dense format during computation.
+
+    Parameters
+    ----------
+    n_features : int
+        Number of input features
+    fit_intercept : bool
+        Whether to fit an intercept term
+    beta : float
+        Temperature parameter for Hard Concrete distribution
+    gamma : float
+        Lower bound of stretched concrete distribution
+    zeta : float
+        Upper bound of stretched concrete distribution
+    init_keep_prob : float
+        Initial probability of keeping each feature
+    device : str or torch.device
+        Device to run computations on ('cpu' or 'cuda')
+    """
+
+    def __init__(
+        self,
+        n_features: int,
+        fit_intercept: bool = True,
+        beta: float = 2 / 3,
+        gamma: float = -0.1,
+        zeta: float = 1.1,
+        init_keep_prob: float = 0.5,
+        device: str | torch.device = "cpu",
+    ):
+        super().__init__()
+        self.n_features = n_features
+        self.fit_intercept = fit_intercept
+        self.beta = beta
+        self.gamma = gamma
+        self.zeta = zeta
+        self.device = torch.device(device)
+
+        # Model parameters
+        self.weight = nn.Parameter(torch.zeros(n_features, device=self.device))
+        if fit_intercept:
+            self.bias = nn.Parameter(torch.zeros(1, device=self.device))
+        else:
+            self.register_parameter("bias", None)
+
+        # L0 gate parameters
+        mu = torch.log(torch.tensor(init_keep_prob / (1 - init_keep_prob)))
+        self.log_alpha = nn.Parameter(
+            torch.normal(
+                mu.item(), 0.01, size=(n_features,), device=self.device
+            )
+        )
+
+        # Cache for sparse tensor conversion
+        self._cached_X_torch: torch.sparse.Tensor | None = None
+        self._cached_X_shape: tuple[int, int] | None = None
+
+    def _convert_sparse_to_torch(
+        self, X_sparse: sp.spmatrix
+    ) -> torch.sparse.Tensor:
+        """
+        Convert scipy sparse matrix to torch sparse tensor.
+
+        Caches the result if the shape matches to avoid redundant conversions.
+        """
+        if (
+            self._cached_X_torch is not None
+            and self._cached_X_shape == X_sparse.shape
+            and X_sparse.nnz == self._cached_X_torch._nnz()
+        ):
+            return self._cached_X_torch
+
+        X_coo = X_sparse.tocoo()
+        indices = torch.LongTensor(np.vstack([X_coo.row, X_coo.col])).to(
+            self.device
+        )
+        values = torch.FloatTensor(X_coo.data).to(self.device)
+        X_torch = torch.sparse_coo_tensor(
+            indices,
+            values,
+            X_sparse.shape,
+            dtype=torch.float32,
+            device=self.device,
+        )
+
+        # Cache for future use
+        self._cached_X_torch = X_torch
+        self._cached_X_shape = X_sparse.shape
+
+        return X_torch
+
+    def _sample_gates(self) -> torch.Tensor:
+        """Sample gates using Hard Concrete distribution."""
+        eps = 1e-6
+        u = torch.rand_like(self.log_alpha).clamp(eps, 1 - eps)
+        X = (torch.log(u) - torch.log(1 - u) + self.log_alpha) / self.beta
+        s = torch.sigmoid(X)
+        s_bar = s * (self.zeta - self.gamma) + self.gamma
+        return s_bar.clamp(0, 1)
+
+    def get_deterministic_gates(self) -> torch.Tensor:
+        """Get deterministic gate values (for inference)."""
+        X = self.log_alpha / self.beta
+        s = torch.sigmoid(X)
+        s_bar = s * (self.zeta - self.gamma) + self.gamma
+        return s_bar.clamp(0, 1)
+
+    def forward(
+        self,
+        X: torch.sparse.Tensor | sp.spmatrix,
+        deterministic: bool = False,
+    ) -> torch.Tensor:
+        """
+        Forward pass maintaining sparse computation.
+
+        Parameters
+        ----------
+        X : torch.sparse.Tensor or scipy.sparse matrix
+            Sparse input tensor of shape (n_samples, n_features)
+        deterministic : bool
+            Whether to use deterministic gates (for inference)
+
+        Returns
+        -------
+        torch.Tensor
+            Predictions of shape (n_samples,)
+        """
+        # Convert scipy sparse to torch if needed
+        if isinstance(X, sp.spmatrix):
+            X = self._convert_sparse_to_torch(X)
+
+        # Sample or get deterministic gates
+        if deterministic:
+            gates = self.get_deterministic_gates()
+        else:
+            gates = self._sample_gates()
+
+        # Apply gates to weights
+        gated_weight = self.weight * gates
+
+        # Sparse matrix multiplication
+        # X is (n_samples, n_features), gated_weight is (n_features,)
+        y = torch.sparse.mm(X, gated_weight.unsqueeze(1)).squeeze(1)
+
+        # Add bias if needed
+        if self.fit_intercept and self.bias is not None:
+            y = y + self.bias
+
+        return y  # type: ignore[no-any-return]
+
+    def get_l0_penalty(self) -> torch.Tensor:
+        """
+        Compute L0 complexity penalty.
+
+        Returns expected number of active parameters.
+        """
+        c = -self.beta * torch.log(
+            torch.tensor(-self.gamma / self.zeta, device=self.device)
+        )
+        pi = torch.sigmoid(self.log_alpha + c)
+        return pi.sum()
+
+    def get_sparsity(self) -> float:
+        """
+        Get current sparsity level.
+
+        Returns
+        -------
+        float
+            Fraction of weights that are effectively zero
+        """
+        with torch.no_grad():
+            gates = self.get_deterministic_gates()
+            return (gates < 0.01).float().mean().item()
+
+    def get_selected_features(self, threshold: float = 0.01) -> torch.Tensor:
+        """
+        Get indices of selected (non-zero) features.
+
+        Parameters
+        ----------
+        threshold : float
+            Gate values below this are considered zero
+
+        Returns
+        -------
+        torch.Tensor
+            Indices of selected features
+        """
+        with torch.no_grad():
+            gates = self.get_deterministic_gates()
+            return torch.where(gates > threshold)[0]
+
+    def fit(
+        self,
+        X_sparse: sp.spmatrix,
+        y: torch.Tensor | np.ndarray,
+        lambda_reg: float = 0.01,
+        lr: float = 0.01,
+        epochs: int = 1000,
+        verbose: bool = False,
+        verbose_freq: int = 100,
+    ) -> "SparseL0Linear":
+        """
+        Fit the model using gradient descent.
+
+        Parameters
+        ----------
+        X_sparse : scipy.sparse matrix
+            Sparse input features
+        y : torch.Tensor or numpy.ndarray
+            Target values
+        lambda_reg : float
+            L0 regularization strength
+        lr : float
+            Learning rate
+        epochs : int
+            Number of training epochs
+        verbose : bool
+            Whether to print progress
+        verbose_freq : int
+            How often to print progress
+
+        Returns
+        -------
+        self
+            Fitted model
+        """
+        import numpy as np
+
+        # Convert y to tensor if needed
+        if isinstance(y, np.ndarray):
+            y = torch.tensor(y, dtype=torch.float32, device=self.device)
+        elif not isinstance(y, torch.Tensor):
+            y = torch.tensor(y, dtype=torch.float32, device=self.device)
+        else:
+            y = y.to(self.device)
+
+        # Convert sparse matrix to torch sparse tensor (will be cached)
+        X_torch = self._convert_sparse_to_torch(X_sparse)
+
+        # Initialize weights with small random values
+        nn.init.normal_(self.weight, 0, 0.1)
+        if self.fit_intercept and self.bias is not None:
+            self.bias.data.fill_(y.mean())
+
+        # Setup optimizer
+        params = [self.weight, self.log_alpha]
+        if self.fit_intercept and self.bias is not None:
+            params.append(self.bias)
+        optimizer = torch.optim.Adam(params, lr=lr)
+
+        # Training loop
+        for epoch in range(epochs):
+            # Forward pass
+            y_pred = self.forward(X_torch, deterministic=False)
+
+            # Compute loss
+            mse_loss = (y - y_pred).pow(2).mean()
+            l0_loss = self.get_l0_penalty()
+            loss = mse_loss + lambda_reg * l0_loss
+
+            # Backward pass
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            # Print progress
+            if verbose and (epoch + 1) % verbose_freq == 0:
+                with torch.no_grad():
+                    n_active = (self.get_deterministic_gates() > 0.01).sum()
+                    print(
+                        f"Epoch {epoch+1:4d}: "
+                        f"loss={loss.item():.4f}, "
+                        f"mse={mse_loss.item():.4f}, "
+                        f"l0={l0_loss.item():.2f}, "
+                        f"active={n_active.item()}"
+                    )
+
+        return self
+
+    def predict(
+        self, X_sparse: sp.spmatrix | torch.sparse.Tensor
+    ) -> torch.Tensor:
+        """
+        Make predictions on new data.
+
+        Parameters
+        ----------
+        X_sparse : scipy.sparse matrix or torch.sparse.Tensor
+            Sparse input features
+
+        Returns
+        -------
+        torch.Tensor
+            Predictions
+        """
+        # Make deterministic predictions
+        with torch.no_grad():
+            return self.forward(X_sparse, deterministic=True)

{l0_python-0.1.0 → l0_python-0.1.1}/l0_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: l0-python
-Version: 0.1.0
+Version: 0.1.1
 Summary: L0 regularization for sparse neural networks and intelligent sampling
 Author-email: PolicyEngine <hello@policyengine.org>
 License: MIT License
@@ -41,6 +41,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0.0
 Requires-Dist: numpy>=1.20.0
+Requires-Dist: scipy>=1.7.0
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov>=3.0; extra == "dev"
@@ -49,6 +50,7 @@ Requires-Dist: mypy>=0.900; extra == "dev"
 Requires-Dist: ruff>=0.0.280; extra == "dev"
 Requires-Dist: pre-commit>=3.0; extra == "dev"
 Requires-Dist: build>=1.3; extra == "dev"
+Requires-Dist: scikit-learn>=1.0; extra == "dev"
 Provides-Extra: docs
 Requires-Dist: mystmd>=1.3; extra == "docs"
 Requires-Dist: jupyter>=1.0; extra == "docs"

{l0_python-0.1.0 → l0_python-0.1.1}/l0_python.egg-info/SOURCES.txt

@@ -6,6 +6,7 @@ l0/distributions.py
 l0/gates.py
 l0/layers.py
 l0/penalties.py
+l0/sparse.py
 l0_python.egg-info/PKG-INFO
 l0_python.egg-info/SOURCES.txt
 l0_python.egg-info/dependency_links.txt
@@ -14,4 +15,5 @@ l0_python.egg-info/top_level.txt
 tests/test_distributions.py
 tests/test_gates.py
 tests/test_layers.py
-tests/test_penalties.py
+tests/test_penalties.py
+tests/test_sparse.py

{l0_python-0.1.0 → l0_python-0.1.1}/l0_python.egg-info/requires.txt

@@ -1,5 +1,6 @@
 torch>=2.0.0
 numpy>=1.20.0
+scipy>=1.7.0
 
 [dev]
 pytest>=7.0
@@ -9,6 +10,7 @@ mypy>=0.900
 ruff>=0.0.280
 pre-commit>=3.0
 build>=1.3
+scikit-learn>=1.0
 
 [docs]
 mystmd>=1.3

{l0_python-0.1.0 → l0_python-0.1.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "l0-python"
-version = "0.1.0"
+version = "0.1.1"
 description = "L0 regularization for sparse neural networks and intelligent sampling"
 readme = "README.md"
 authors = [
@@ -23,6 +23,7 @@ classifiers = [
 dependencies = [
     "torch>=2.0.0",
     "numpy>=1.20.0",
+    "scipy>=1.7.0",
 ]
 keywords = ["l0", "regularization", "sparsity", "neural networks", "calibration"]
 
@@ -35,6 +36,7 @@ dev = [
     "ruff>=0.0.280",
     "pre-commit>=3.0",
     "build>=1.3",
+    "scikit-learn>=1.0",
 ]
 docs = [
     "mystmd>=1.3",

l0_python-0.1.1/tests/test_sparse.py

@@ -0,0 +1,240 @@
+"""
+Tests for sparse L0 linear models.
+"""
+
+# ruff: noqa: N806
+
+import numpy as np
+import pytest
+import torch
+from scipy import sparse as sp
+from sklearn.linear_model import LinearRegression
+
+from l0.sparse import SparseL0Linear
+
+
+class TestSparseL0Linear:
+    """Test cases for SparseL0Linear model."""
+
+    def test_sparse_data_recovery(self):
+        """
+        Test that the model can recover sparse coefficients from sparse data.
+        Uses the same data generating process as the demonstration.
+        """
+        # Set seeds for reproducibility
+        torch.manual_seed(12543)
+        np.random.seed(12543)
+
+        # Data parameters
+        n = 1000
+        p = 100
+        X_sparsity = 0.95  # 95% of X entries are zero
+        beta_sparsity = 0.95  # 95% of true coefficients are zero
+
+        # Create sparse true coefficients
+        num_nonzero = int(p * (1 - beta_sparsity))
+        nonzero_indices = np.random.choice(p, num_nonzero, replace=False)
+        b = np.zeros(p)
+        b[nonzero_indices] = np.random.choice(
+            [-3, -2, -1, 1, 2, 3], num_nonzero
+        )
+
+        # Generate sparse X matrix
+        X_dense = np.random.randn(n, p) * 0.5
+        mask = np.random.random((n, p)) < X_sparsity
+        X_dense[mask] = 0
+
+        # Add correlation structure
+        for i in range(0, p - 1, 2):
+            if i + 1 < p:
+                corr_mask = ~(mask[:, i] | mask[:, i + 1])
+                X_dense[corr_mask, i + 1] += 0.3 * X_dense[corr_mask, i]
+
+        # Convert to sparse
+        X_sparse = sp.csr_matrix(X_dense)
+
+        # Generate y with intercept
+        b0 = 30
+        sigma_e = 1.5
+        y_true = b0 + X_sparse @ b
+        y = y_true + sigma_e * np.random.randn(n)
+
+        # Fit baseline linear regression for lambda selection
+        lr = LinearRegression(fit_intercept=True).fit(X_sparse, y)
+        residuals = y - lr.predict(X_sparse)
+        sigma2_hat = np.var(residuals)
+        lambda_reg = 0.01 * sigma2_hat
+
+        # Fit our sparse L0 model
+        model = SparseL0Linear(
+            n_features=p, fit_intercept=True, init_keep_prob=0.5
+        )
+
+        model.fit(
+            X_sparse,
+            y,
+            lambda_reg=lambda_reg,
+            lr=0.01,
+            epochs=2000,
+            verbose=False,
+        )
+
+        # Check feature selection performance
+        selected = model.get_selected_features().cpu().numpy()
+        true_positives = len(set(selected) & set(nonzero_indices))
+
+        # We expect to recover at least 60% of true features
+        recall = (
+            true_positives / len(nonzero_indices)
+            if len(nonzero_indices) > 0
+            else 0
+        )
+        assert recall >= 0.6, f"Recall {recall:.2%} is too low"
+
+        # We expect reasonable precision (at least 50%)
+        if len(selected) > 0:
+            precision = true_positives / len(selected)
+            assert precision >= 0.5, f"Precision {precision:.2%} is too low"
+
+        # Check that sparsity is achieved
+        sparsity = model.get_sparsity()
+        assert sparsity >= 0.7, f"Sparsity {sparsity:.2%} is too low"
+
+        # Check prediction quality on training data
+        y_pred = model.predict(X_sparse).cpu().numpy()
+        mse = np.mean((y - y_pred) ** 2)
+        baseline_mse = np.mean((y - lr.predict(X_sparse)) ** 2)
+
+        # Should be comparable to baseline
+        assert mse < 1.5 * baseline_mse, "MSE is much worse than baseline"
+
+    def test_no_intercept(self):
+        """Test that the model works without intercept."""
+        np.random.seed(42)
+        torch.manual_seed(42)
+
+        n, p = 100, 20
+        X_dense = np.random.randn(n, p)
+        mask = np.random.random((n, p)) < 0.8
+        X_dense[mask] = 0
+        X_sparse = sp.csr_matrix(X_dense)
+
+        # Generate y without intercept
+        b = np.zeros(p)
+        b[:3] = [2, -3, 1]
+        y = X_sparse @ b + 0.1 * np.random.randn(n)
+
+        # Fit model without intercept
+        model = SparseL0Linear(
+            n_features=p, fit_intercept=False, init_keep_prob=0.5
+        )
+
+        assert (
+            model.bias is None
+        ), "Bias should be None when fit_intercept=False"
+
+        model.fit(X_sparse, y, epochs=500, verbose=False)
+
+        # Check that we find the right features
+        selected = model.get_selected_features().cpu().numpy()
+        assert (
+            0 in selected or 1 in selected or 2 in selected
+        ), "Should select at least one true feature"
+
+        # Make predictions
+        y_pred = model.predict(X_sparse)
+        assert y_pred.shape == (n,), "Predictions should have correct shape"
+
+    def test_sparse_efficiency(self):
+        """Test that sparse computations remain sparse."""
+        np.random.seed(42)
+        torch.manual_seed(42)
+
+        # Create very sparse data
+        n, p = 500, 1000
+        X_dense = np.random.randn(n, p)
+        mask = np.random.random((n, p)) < 0.99  # 99% sparse
+        X_dense[mask] = 0
+        X_sparse = sp.csr_matrix(X_dense)
+
+        # Check sparsity
+        sparsity = 1 - X_sparse.nnz / (n * p)
+        assert sparsity > 0.98, "Data should be very sparse"
+
+        y = np.random.randn(n)
+
+        model = SparseL0Linear(n_features=p, fit_intercept=True)
+
+        # The model should handle this efficiently
+        model.fit(X_sparse, y, epochs=100, verbose=False)
+
+        # Check that cached tensor is sparse
+        assert (
+            model._cached_X_torch is not None
+        ), "Should cache the sparse tensor"
+        assert (
+            model._cached_X_torch.is_sparse
+        ), "Cached tensor should be sparse"
+
+        # Multiple predictions should reuse cached tensor
+        y_pred1 = model.predict(X_sparse)
+        y_pred2 = model.predict(X_sparse)
+
+        # Should get same predictions
+        torch.testing.assert_close(y_pred1, y_pred2)
+
+    def test_device_support(self):
+        """Test that model works on different devices."""
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA not available")
+
+        np.random.seed(42)
+        torch.manual_seed(42)
+
+        n, p = 50, 10
+        X_dense = np.random.randn(n, p)
+        X_sparse = sp.csr_matrix(X_dense)
+        y = np.random.randn(n)
+
+        # Test on CUDA
+        model = SparseL0Linear(n_features=p, fit_intercept=True, device="cuda")
+
+        assert model.weight.device.type == "cuda"
+        assert model.log_alpha.device.type == "cuda"
+
+        model.fit(X_sparse, y, epochs=100, verbose=False)
+        y_pred = model.predict(X_sparse)
+
+        assert y_pred.device.type == "cuda"
+
+    def test_deterministic_vs_stochastic(self):
+        """Test that deterministic mode gives consistent predictions."""
+        np.random.seed(42)
+        torch.manual_seed(42)
+
+        n, p = 100, 20
+        X_dense = np.random.randn(n, p)
+        X_sparse = sp.csr_matrix(X_dense)
+        y = np.random.randn(n)
+
+        model = SparseL0Linear(n_features=p, fit_intercept=True)
+        model.fit(X_sparse, y, epochs=200, verbose=False)
+
+        # Convert X for forward pass
+        X_torch = model._convert_sparse_to_torch(X_sparse)
+
+        # Deterministic predictions should be the same
+        with torch.no_grad():
+            y_det1 = model.forward(X_torch, deterministic=True)
+            y_det2 = model.forward(X_torch, deterministic=True)
+
+        torch.testing.assert_close(y_det1, y_det2)
+
+        # Stochastic predictions should differ
+        with torch.no_grad():
+            y_stoch1 = model.forward(X_torch, deterministic=False)
+            y_stoch2 = model.forward(X_torch, deterministic=False)
+
+        assert not torch.allclose(
+            y_stoch1, y_stoch2
+        ), "Stochastic predictions should differ"