l0-python 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

l0/__init__.py

@@ -5,7 +5,7 @@ A PyTorch implementation of L0 regularization for neural network sparsification
 and intelligent sampling, based on Louizos, Welling, & Kingma (2017).
 """
 
-__version__ = "0.1.0"
+__version__ = "0.5.0"
 
 from .distributions import HardConcrete
 from .gates import FeatureGate, HybridGate, L0Gate, SampleGate

l0/calibration.py

@@ -60,6 +60,7 @@ class SparseCalibrationWeights(nn.Module):
         log_weight_jitter_sd: float = 0.0,
         log_alpha_jitter_sd: float = 0.01,
         device: str | torch.device = "cpu",
+        use_gates: bool = True,
     ):
         super().__init__()
         self.n_features = n_features
@@ -69,6 +70,7 @@ class SparseCalibrationWeights(nn.Module):
         self.log_weight_jitter_sd = log_weight_jitter_sd
         self.log_alpha_jitter_sd = log_alpha_jitter_sd
         self.device = torch.device(device)
+        self.use_gates = use_gates
 
         # Initialize weights (on original scale)
         if init_weights is None:
@@ -193,6 +195,10 @@ class SparseCalibrationWeights(nn.Module):
         torch.Tensor
             Positive calibration weights with L0 sparsity applied
         """
+        if not self.use_gates:
+            # No gates - just return positive weights
+            return torch.exp(self.log_weight)
+
         # Sample or get deterministic gates
         if deterministic:
             gates = self.get_deterministic_gates()
@@ -453,6 +459,46 @@ class SparseCalibrationWeights(nn.Module):
                         1 - active_info["count"] / self.n_features
                     )
 
+                    # Weight distribution diagnostic
+                    if active_weights.numel() > 0:
+                        w_tiny = (active_weights < 0.01).sum().item()
+                        w_small = (
+                            ((active_weights >= 0.01) & (active_weights < 0.1))
+                            .sum()
+                            .item()
+                        )
+                        w_med = (
+                            ((active_weights >= 0.1) & (active_weights < 1.0))
+                            .sum()
+                            .item()
+                        )
+                        w_normal = (
+                            ((active_weights >= 1.0) & (active_weights < 10.0))
+                            .sum()
+                            .item()
+                        )
+                        w_large = (
+                            (
+                                (active_weights >= 10.0)
+                                & (active_weights < 1000.0)
+                            )
+                            .sum()
+                            .item()
+                        )
+                        w_huge = (active_weights >= 1000.0).sum().item()
+
+                        total_active = active_weights.numel()
+                        weight_dist = (
+                            f"[<0.01: {100*w_tiny/total_active:.1f}%, "
+                            f"0.01-0.1: {100*w_small/total_active:.1f}%, "
+                            f"0.1-1: {100*w_med/total_active:.1f}%, "
+                            f"1-10: {100*w_normal/total_active:.1f}%, "
+                            f"10-1000: {100*w_large/total_active:.1f}%, "
+                            f">1000: {100*w_huge/total_active:.1f}%]"
+                        )
+                    else:
+                        weight_dist = "[no active weights]"
+
                     # Calculate components of the actual loss being minimized
                     actual_data_loss = data_loss.item()
                     actual_l0_loss = l0_loss.item()
@@ -464,7 +510,8 @@ class SparseCalibrationWeights(nn.Module):
                             f"mean_group_mare={mean_group_mare:.4%}, "
                             f"max_error={max_rel_err:.1%}, "
                             f"total_loss={actual_total_loss:.3f}, "
-                            f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)"
+                            f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)\n"
+                            f"         Weight dist: {weight_dist}"
                         )
                     else:
                         print(
@@ -472,7 +519,8 @@ class SparseCalibrationWeights(nn.Module):
                             f"mean_error={mean_rel_err:.4%}, "
                             f"max_error={max_rel_err:.1%}, "
                             f"total_loss={actual_total_loss:.3f}, "
-                            f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)"
+                            f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)\n"
+                            f"         Weight dist: {weight_dist}"
                         )
 
         return self

{l0_python-0.4.2.dist-info → l0_python-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: l0-python
-Version: 0.4.2
+Version: 0.5.0
 Summary: L0 regularization for sparse neural networks and intelligent sampling
 Author-email: PolicyEngine <hello@policyengine.org>
 License: MIT License
@@ -68,21 +68,7 @@ Dynamic: license-file
 [![PyPI version](https://badge.fury.io/py/l0-python.svg)](https://pypi.org/project/l0-python/)
 [![CI](https://github.com/PolicyEngine/L0/actions/workflows/push.yml/badge.svg)](https://github.com/PolicyEngine/L0/actions)
 
-A PyTorch implementation of L0 regularization based on [Louizos, Welling, & Kingma (2017)](https://arxiv.org/abs/1712.01312), with a **critical bug fix** to the original authors' implementation.
-
-## Why This Package?
-
-The original L0 implementation from AMLab-Amsterdam contains a bug in test-time gate computation where the temperature parameter is incorrectly omitted. This prevents proper sparsity and degrades performance. Our implementation corrects this error:
-
-```python
-# Original (buggy): gates never fully close
-pi = sigmoid(log_alpha)
-
-# Corrected: temperature division required for proper sparsity
-pi = sigmoid(log_alpha / temperature)
-```
-
-This fix enables gates to achieve true 0/1 values, producing exact sparsity as intended by the L0 formulation.
+A PyTorch implementation of L0 regularization based on [Louizos, Welling, & Kingma (2017)](https://arxiv.org/abs/1712.01312), designed for survey calibration and sparse regression.
 
 ## Installation
 
@@ -97,9 +83,24 @@ cd L0
 pip install -e .[dev]
 ```
 
+## Our Approach to Test-Time Gates
+
+The original Hard Concrete formulation uses temperature (β) during training to control the sharpness of stochastic gates. At test time, there's a design choice: whether to include temperature in the deterministic gate computation.
+
+We include temperature at test time:
+```python
+# Our approach: include temperature
+z = sigmoid(log_alpha / beta) * (zeta - gamma) + gamma
+
+# Alternative: omit temperature
+z = sigmoid(log_alpha) * (zeta - gamma) + gamma
+```
+
+Including temperature produces sharper 0/1 decisions, which we find beneficial for achieving clean sparsity in our applications. See `examples/sparse_regression_demo.py` for a demonstration on a 4-variable regression problem.
+
 ## Primary Use Case: Survey Calibration
 
-This package was developed for PolicyEngine's survey calibration system, where we select a sparse subset of survey households while matching population targets.
+This package was developed for PolicyEngine's survey calibration, where we select a sparse subset of survey households while matching population targets.
 
 ```python
 import numpy as np
@@ -117,119 +118,73 @@ model = SparseCalibrationWeights(
     beta=0.35,
     gamma=-0.1,
     zeta=1.1,
-    init_keep_prob=0.5,      # Start with 50% active probability
-    init_weights=1.0,        # Or pass array of initial weights
+    init_keep_prob=0.5,
+    init_weights=1.0,
     log_weight_jitter_sd=0.05,
-    device="cuda",           # GPU acceleration
+    device="cuda",
 )
 
 # Train with L0+L2 regularization
 model.fit(
     M=M,
     y=y,
-    lambda_l0=1e-6,          # Controls sparsity level
-    lambda_l2=1e-8,          # Prevents weight explosion
+    lambda_l0=1e-6,
+    lambda_l2=1e-8,
     lr=0.15,
     epochs=2000,
-    loss_type="relative",    # Scale-invariant loss
+    loss_type="relative",
     verbose=True,
-    verbose_freq=100,
 )
 
 # Get results
 active = model.get_active_weights()
-print(f"Selected {active['count']} of {N} households ({100*active['count']/N:.1f}%)")
+print(f"Selected {active['count']} of {N} households")
 print(f"Sparsity: {model.get_sparsity():.1%}")
-
-# Predict calibrated totals
-y_pred = model.predict(M)
 ```
 
-### Key Features for Calibration
+### Key Features
 
-- **Non-negative weights**: All weights constrained to be positive via log-space parameterization
-- **Sparse solutions**: L0 penalty directly minimizes the count of active weights
-- **Relative loss**: Scale-invariant loss for targets spanning orders of magnitude
-- **Group-wise averaging**: Balance loss contributions across target groups with different cardinalities
-- **GPU support**: CUDA acceleration for large-scale problems
+- **Non-negative weights**: Constrained via log-space parameterization
+- **L0 sparsity**: Directly minimizes the count of active weights
+- **Relative loss**: Scale-invariant for targets spanning orders of magnitude
+- **Group-wise averaging**: Balance loss across target groups with different sizes
+- **GPU support**: CUDA acceleration for large problems
 
-## Neural Network Sparsification
+## Sparse Regression
 
-The package also supports traditional neural network pruning:
+For sparse linear regression with scipy sparse matrices:
 
 ```python
-import torch
-from l0 import L0Linear, compute_l0l2_penalty, TemperatureScheduler, update_temperatures
-
-class SparseModel(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.fc1 = L0Linear(784, 256, init_sparsity=0.5)
-        self.fc2 = L0Linear(256, 10, init_sparsity=0.7)
-
-    def forward(self, x):
-        x = torch.relu(self.fc1(x))
-        return self.fc2(x)
-
-model = SparseModel()
-optimizer = torch.optim.Adam(model.parameters())
-scheduler = TemperatureScheduler(initial_temp=1.0, final_temp=0.1)
-
-for epoch in range(100):
-    temp = scheduler.get_temperature(epoch)
-    update_temperatures(model, temp)
-
-    output = model(input_data)
-    ce_loss = criterion(output, target)
-    penalty = compute_l0l2_penalty(model, l0_lambda=1e-3, l2_lambda=1e-4)
-    loss = ce_loss + penalty
-
-    optimizer.zero_grad()
-    loss.backward()
-    optimizer.step()
-```
-
-### Available Layers
-
-- `L0Linear`: Fully connected layer with L0 gates on weights
-- `L0Conv2d`: 2D convolution with channel-wise L0 gates
-- `L0DepthwiseConv2d`: Depthwise convolution with L0 gates
-- `SparseMLP`: Multi-layer perceptron with built-in L0 regularization
+from scipy import sparse as sp
+from l0.sparse import SparseL0Linear
 
-## Intelligent Sampling Gates
+# Sparse design matrix
+X = sp.random(1000, 500, density=0.1, format="csr")
+y = np.random.randn(1000)
 
-Standalone gates for sample/feature selection:
+model = SparseL0Linear(n_features=500)
+model.fit(X, y, lambda_l0=0.001, epochs=1000)
 
-```python
-from l0 import SampleGate, FeatureGate, HybridGate
+# Get sparse coefficients
+coef = model.get_coefficients(threshold=0.01)
+```
 
-# Select samples via learned gates
-gate = SampleGate(n_samples=10000, target_samples=1000)
-selected_data, indices = gate.select_samples(data)
+## Example: Variable Selection
 
-# Select features
-gate = FeatureGate(n_features=1000, max_features=50)
-selected_data, indices = gate.select_features(data)
+The `examples/sparse_regression_demo.py` script demonstrates L0 regularization on a simple problem where the true coefficients are `[1, 0, -2, 0]`:
 
-# Hybrid: combine L0 selection with random sampling
-hybrid = HybridGate(
-    n_items=10000,
-    l0_fraction=0.25,    # 25% via learned L0 gates
-    random_fraction=0.75 # 75% random for coverage
-)
-selected, indices, types = hybrid.select(data)
+```bash
+python examples/sparse_regression_demo.py
 ```
 
-## How L0 Regularization Works
-
-Unlike post-hoc pruning (setting small weights to zero), L0 regularization integrates sparsity into the training objective:
-
-1. **Stochastic gates**: Each weight has a learned gate parameter controlling activation probability
-2. **Hard Concrete distribution**: Enables differentiable 0/1 gate values during training
-3. **Expected L0 penalty**: Minimizes the expected number of active gates
-4. **Temperature annealing**: Gradually sharpens gates from soft to hard decisions
+Output:
+```
+True coefficients:        [ 1.  0. -2.  0.]
+Recovered coefficients:   [ 1.039  0.    -2.069 -0.   ]
+Gates:                    [1. 0. 1. 0.]
+```
 
-The result: the network learns *which* weights should be zero as part of optimization, not as a post-processing step.
+The model correctly identifies that only variables 1 and 3 contribute to the outcome.
 
 ## Testing
 
@@ -239,8 +194,6 @@ pytest tests/ -v --cov=l0
 
 ## Citation
 
-If you use this package, please cite the original paper:
-
 ```bibtex
 @article{louizos2017learning,
   title={Learning Sparse Neural Networks through L0 Regularization},

l0_python-0.5.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+l0/__init__.py,sha256=ImRgZfVg-RJhLoqNaLc8ysjddZ7CYGHJPJbRUyfSoHY,1192
+l0/calibration.py,sha256=1-C6B6O47l6qlEcQl3N8waAPAyj5PFlzBHl17KbXqGo,20018
+l0/distributions.py,sha256=GvVoDasdxoDlPPdzt33lFzkuJgQsfkEn9e-abxru3nM,6325
+l0/gates.py,sha256=54OORMg9qSsdbIyzYIL9t9CCHWsKbumSKl_d4-YFuNU,11443
+l0/layers.py,sha256=j_DPEFdfUc37ojtoK2ufGzVYzZ_DwusqTcJXCVxZS68,13785
+l0/penalties.py,sha256=okqXgX2eGUmg5prcDclicTQALPPPF9FA5eQP09V4GTE,9569
+l0/sparse.py,sha256=UVpkrrJBnQ9iSNDoWK5F_sr43tQareWXzYHyMq_cKss,9726
+l0_python-0.5.0.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
+l0_python-0.5.0.dist-info/METADATA,sha256=zoUDdaSbSAWeWtWW-K_5Wgl6RbTqGHfGIyfYobvxEBQ,7095
+l0_python-0.5.0.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+l0_python-0.5.0.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
+l0_python-0.5.0.dist-info/RECORD,,

{l0_python-0.4.2.dist-info → l0_python-0.5.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

l0_python-0.4.2.dist-info/RECORD

@@ -1,12 +0,0 @@
-l0/__init__.py,sha256=ezGVLGbdU_jAElAJSdLKjSWCnZRILAZi8qMBPjcHwQc,1192
-l0/calibration.py,sha256=3VM-28M1yp-DM6AlHX7sqZCaE7OJevpY9Ey-syjSPik,17838
-l0/distributions.py,sha256=GvVoDasdxoDlPPdzt33lFzkuJgQsfkEn9e-abxru3nM,6325
-l0/gates.py,sha256=54OORMg9qSsdbIyzYIL9t9CCHWsKbumSKl_d4-YFuNU,11443
-l0/layers.py,sha256=j_DPEFdfUc37ojtoK2ufGzVYzZ_DwusqTcJXCVxZS68,13785
-l0/penalties.py,sha256=okqXgX2eGUmg5prcDclicTQALPPPF9FA5eQP09V4GTE,9569
-l0/sparse.py,sha256=UVpkrrJBnQ9iSNDoWK5F_sr43tQareWXzYHyMq_cKss,9726
-l0_python-0.4.2.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
-l0_python-0.4.2.dist-info/METADATA,sha256=czuRwTuF5iOn6f-rsFSxk_gHXWuKXAaYAmDugkGw4rA,9074
-l0_python-0.4.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-l0_python-0.4.2.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
-l0_python-0.4.2.dist-info/RECORD,,