l0-python 0.1.0__tar.gz

l0_python-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 PolicyEngine
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

l0_python-0.1.0/PKG-INFO

@@ -0,0 +1,215 @@
+Metadata-Version: 2.4
+Name: l0-python
+Version: 0.1.0
+Summary: L0 regularization for sparse neural networks and intelligent sampling
+Author-email: PolicyEngine <hello@policyengine.org>
+License: MIT License
+        
+        Copyright (c) 2025 PolicyEngine
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Homepage, https://github.com/PolicyEngine/L0
+Project-URL: Repository, https://github.com/PolicyEngine/L0
+Project-URL: Documentation, https://github.com/PolicyEngine/L0#readme
+Project-URL: Issues, https://github.com/PolicyEngine/L0/issues
+Keywords: l0,regularization,sparsity,neural networks,calibration
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0.0
+Requires-Dist: numpy>=1.20.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov>=3.0; extra == "dev"
+Requires-Dist: black>=22.0; extra == "dev"
+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"
+Provides-Extra: docs
+Requires-Dist: mystmd>=1.3; extra == "docs"
+Requires-Dist: jupyter>=1.0; extra == "docs"
+Requires-Dist: notebook>=7.0; extra == "docs"
+Requires-Dist: ipykernel>=6.0; extra == "docs"
+Provides-Extra: examples
+Requires-Dist: matplotlib>=3.3; extra == "examples"
+Requires-Dist: tqdm>=4.60; extra == "examples"
+Requires-Dist: pandas>=1.3; extra == "examples"
+Requires-Dist: scikit-learn>=1.0; extra == "examples"
+Requires-Dist: torchvision>=0.15; extra == "examples"
+Dynamic: license-file
+
+# L0 Regularization
+
+A PyTorch implementation of L0 regularization for neural network sparsification and intelligent sampling, based on [Louizos, Welling, & Kingma (2017)](https://arxiv.org/abs/1712.01312).
+
+## Features
+
+- **Hard Concrete Distribution**: Differentiable approximation of L0 norm
+- **Sparse Neural Network Layers**: L0Linear, L0Conv2d with automatic pruning
+- **Intelligent Sampling**: Sample/feature selection gates for calibration
+- **L0L2 Combined Penalty**: Recommended approach to prevent overfitting
+- **Temperature Scheduling**: Annealing for improved convergence
+- **TDD Development**: Comprehensive test coverage
+
+## Installation
+
+```bash
+pip install l0
+```
+
+For development:
+```bash
+git clone https://github.com/PolicyEngine/L0.git
+cd L0
+pip install -e .[dev]
+```
+
+## Quick Start
+
+### Neural Network Sparsification
+
+```python
+import torch
+from l0 import L0Linear, compute_l0l2_penalty, TemperatureScheduler, update_temperatures
+
+# Create a sparse model
+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)
+
+# Training loop
+for epoch in range(100):
+    # Update temperature
+    temp = scheduler.get_temperature(epoch)
+    update_temperatures(model, temp)
+    
+    # Forward pass
+    output = model(input_data)
+    ce_loss = criterion(output, target)
+    
+    # Add L0L2 penalty
+    penalty = compute_l0l2_penalty(model, l0_lambda=1e-3, l2_lambda=1e-4)
+    loss = ce_loss + penalty
+    
+    # Backward pass
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+```
+
+### Intelligent Sample Selection
+
+```python
+from l0 import SampleGate, HybridGate
+
+# Pure L0 selection
+gate = SampleGate(n_samples=10000, target_samples=1000)
+selected_data, indices = gate.select_samples(data)
+
+# Hybrid selection (25% L0, 75% random)
+hybrid = HybridGate(
+    n_items=10000,
+    l0_fraction=0.25,
+    random_fraction=0.75,
+    target_items=1000
+)
+selected, indices, types = hybrid.select(data)
+```
+
+### Feature Selection
+
+```python
+from l0 import FeatureGate
+
+# Select top features
+gate = FeatureGate(n_features=1000, max_features=50)
+selected_data, feature_indices = gate.select_features(data)
+
+# Get feature importance
+importance = gate.get_feature_importance()
+```
+
+## Integration with PolicyEngine
+
+This package is designed to work with PolicyEngine's calibration system:
+
+```python
+# In policyengine-us-data or similar
+from l0 import HardConcrete
+
+# Use for household selection in CPS calibration
+gates = HardConcrete(
+    len(household_weights),
+    temperature=0.25,
+    init_mean=0.999  # Start with most households
+)
+
+# Apply gates during reweighting
+masked_weights = weights * gates()
+```
+
+## Documentation
+
+Full documentation available at: https://policyengine.github.io/L0/
+
+## Testing
+
+Run tests with:
+```bash
+pytest tests/ -v --cov=l0
+```
+
+## Acknowledgments
+
+This implementation is inspired by and builds upon the [original L0 regularization code](https://github.com/AMLab-Amsterdam/L0_regularization) by AMLab Amsterdam, which accompanied the paper by Louizos et al. (2018).
+
+## Citation
+
+If you use this package, please cite:
+
+```bibtex
+@article{louizos2017learning,
+  title={Learning Sparse Neural Networks through L0 Regularization},
+  author={Louizos, Christos and Welling, Max and Kingma, Diederik P},
+  journal={arXiv preprint arXiv:1712.01312},
+  year={2017}
+}
+```
+
+## License
+
+MIT License - see [LICENSE](LICENSE) file for details.

l0_python-0.1.0/README.md

@@ -0,0 +1,151 @@
+# L0 Regularization
+
+A PyTorch implementation of L0 regularization for neural network sparsification and intelligent sampling, based on [Louizos, Welling, & Kingma (2017)](https://arxiv.org/abs/1712.01312).
+
+## Features
+
+- **Hard Concrete Distribution**: Differentiable approximation of L0 norm
+- **Sparse Neural Network Layers**: L0Linear, L0Conv2d with automatic pruning
+- **Intelligent Sampling**: Sample/feature selection gates for calibration
+- **L0L2 Combined Penalty**: Recommended approach to prevent overfitting
+- **Temperature Scheduling**: Annealing for improved convergence
+- **TDD Development**: Comprehensive test coverage
+
+## Installation
+
+```bash
+pip install l0
+```
+
+For development:
+```bash
+git clone https://github.com/PolicyEngine/L0.git
+cd L0
+pip install -e .[dev]
+```
+
+## Quick Start
+
+### Neural Network Sparsification
+
+```python
+import torch
+from l0 import L0Linear, compute_l0l2_penalty, TemperatureScheduler, update_temperatures
+
+# Create a sparse model
+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)
+
+# Training loop
+for epoch in range(100):
+    # Update temperature
+    temp = scheduler.get_temperature(epoch)
+    update_temperatures(model, temp)
+    
+    # Forward pass
+    output = model(input_data)
+    ce_loss = criterion(output, target)
+    
+    # Add L0L2 penalty
+    penalty = compute_l0l2_penalty(model, l0_lambda=1e-3, l2_lambda=1e-4)
+    loss = ce_loss + penalty
+    
+    # Backward pass
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+```
+
+### Intelligent Sample Selection
+
+```python
+from l0 import SampleGate, HybridGate
+
+# Pure L0 selection
+gate = SampleGate(n_samples=10000, target_samples=1000)
+selected_data, indices = gate.select_samples(data)
+
+# Hybrid selection (25% L0, 75% random)
+hybrid = HybridGate(
+    n_items=10000,
+    l0_fraction=0.25,
+    random_fraction=0.75,
+    target_items=1000
+)
+selected, indices, types = hybrid.select(data)
+```
+
+### Feature Selection
+
+```python
+from l0 import FeatureGate
+
+# Select top features
+gate = FeatureGate(n_features=1000, max_features=50)
+selected_data, feature_indices = gate.select_features(data)
+
+# Get feature importance
+importance = gate.get_feature_importance()
+```
+
+## Integration with PolicyEngine
+
+This package is designed to work with PolicyEngine's calibration system:
+
+```python
+# In policyengine-us-data or similar
+from l0 import HardConcrete
+
+# Use for household selection in CPS calibration
+gates = HardConcrete(
+    len(household_weights),
+    temperature=0.25,
+    init_mean=0.999  # Start with most households
+)
+
+# Apply gates during reweighting
+masked_weights = weights * gates()
+```
+
+## Documentation
+
+Full documentation available at: https://policyengine.github.io/L0/
+
+## Testing
+
+Run tests with:
+```bash
+pytest tests/ -v --cov=l0
+```
+
+## Acknowledgments
+
+This implementation is inspired by and builds upon the [original L0 regularization code](https://github.com/AMLab-Amsterdam/L0_regularization) by AMLab Amsterdam, which accompanied the paper by Louizos et al. (2018).
+
+## Citation
+
+If you use this package, please cite:
+
+```bibtex
+@article{louizos2017learning,
+  title={Learning Sparse Neural Networks through L0 Regularization},
+  author={Louizos, Christos and Welling, Max and Kingma, Diederik P},
+  journal={arXiv preprint arXiv:1712.01312},
+  year={2017}
+}
+```
+
+## License
+
+MIT License - see [LICENSE](LICENSE) file for details.

l0_python-0.1.0/l0/__init__.py

@@ -0,0 +1,53 @@
+"""
+L0 Regularization Package
+
+A PyTorch implementation of L0 regularization for neural network sparsification
+and intelligent sampling, based on Louizos, Welling, & Kingma (2017).
+"""
+
+__version__ = "0.1.0"
+
+from .distributions import HardConcrete
+from .gates import FeatureGate, HybridGate, L0Gate, SampleGate
+from .layers import (
+    L0Conv2d,
+    L0DepthwiseConv2d,
+    L0Linear,
+    SparseMLP,
+    prune_model,
+)
+from .penalties import (
+    PenaltyTracker,
+    TemperatureScheduler,
+    compute_l0_penalty,
+    compute_l0l2_penalty,
+    compute_l2_penalty,
+    get_active_parameter_count,
+    get_sparsity_stats,
+    update_temperatures,
+)
+
+__all__ = [
+    # Distributions
+    "HardConcrete",
+    # Layers
+    "L0Linear",
+    "L0Conv2d",
+    "L0DepthwiseConv2d",
+    "SparseMLP",
+    "prune_model",
+    # Gates
+    "L0Gate",
+    "SampleGate",
+    "FeatureGate",
+    "HybridGate",
+    # Penalties
+    "compute_l0_penalty",
+    "compute_l2_penalty",
+    "compute_l0l2_penalty",
+    "get_sparsity_stats",
+    "get_active_parameter_count",
+    "TemperatureScheduler",
+    "update_temperatures",
+    "PenaltyTracker",
+]

l0_python-0.1.0/l0/distributions.py

@@ -0,0 +1,205 @@
+"""
+Hard Concrete distribution for L0 regularization.
+
+Based on Louizos, Welling, & Kingma (2017): "Learning Sparse Neural Networks
+through L0 Regularization" (https://arxiv.org/abs/1712.01312)
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class HardConcrete(nn.Module):
+    """
+    Hard Concrete distribution for L0 regularization.
+
+    This distribution enables differentiable approximation of the L0 norm
+    by using stochastic gates that can be optimized via gradient descent.
+
+    Parameters
+    ----------
+    input_dim : int
+        Size of the first dimension (e.g., number of samples or features)
+    output_dim : Optional[int]
+        Size of the second dimension for 2D gates (e.g., for weight matrices)
+    temperature : float
+        Temperature parameter controlling the hardness of the distribution.
+        Lower values make the distribution more discrete.
+    stretch : float
+        Stretch parameter for the hard concrete transformation
+    init_mean : float
+        Initial mean activation probability for the gates (0 to 1)
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int | None = None,
+        temperature: float = 0.5,
+        stretch: float = 0.1,
+        init_mean: float = 0.5,
+    ):
+        super().__init__()
+
+        # Set gate dimensions
+        if output_dim is None:
+            self.gate_size: tuple[int, ...] = (input_dim,)
+        else:
+            self.gate_size = (input_dim, output_dim)
+
+        # Learnable parameters - logits for each gate
+        self.qz_logits = nn.Parameter(torch.zeros(self.gate_size))
+
+        # Distribution parameters
+        self.temperature = temperature
+        self.stretch = stretch
+        self.gamma = -stretch  # Lower bound after stretching
+        self.zeta = 1.0 + stretch  # Upper bound after stretching
+        self.init_mean = init_mean
+
+        # Initialize parameters
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """Initialize gate logits based on desired initial mean."""
+        if self.init_mean is not None:
+            # Convert mean probability to logit space
+            init_val = math.log(self.init_mean / (1 - self.init_mean))
+            self.qz_logits.data.fill_(init_val)
+
+    def forward(
+        self, input_shape: tuple[int, ...] | None = None
+    ) -> torch.Tensor:
+        """
+        Sample or compute gates.
+
+        Parameters
+        ----------
+        input_shape : Optional[Tuple[int, ...]]
+            Shape of the input tensor for broadcasting gates if needed
+
+        Returns
+        -------
+        torch.Tensor
+            Gate values in [0, 1], either sampled (training) or deterministic (eval)
+        """
+        if self.training:
+            gates = self._sample_gates()
+        else:
+            gates = self._deterministic_gates()
+
+        # Broadcast gates if needed for convolution layers
+        if input_shape is not None and len(input_shape) > len(gates.shape):
+            # Add dimensions for spatial dims (e.g., height, width in Conv2d)
+            for _ in range(len(input_shape) - len(gates.shape)):
+                gates = gates.unsqueeze(-1)
+
+        return gates
+
+    def _sample_gates(self) -> torch.Tensor:
+        """
+        Sample gates using the reparameterization trick.
+
+        Returns
+        -------
+        torch.Tensor
+            Sampled gate values in [0, 1]
+        """
+        # Sample uniform noise (avoiding exact 0 and 1 for numerical stability)
+        u = torch.zeros_like(self.qz_logits).uniform_(1e-8, 1.0 - 1e-8)
+
+        # Apply the concrete distribution transformation
+        # s = sigmoid((log(u) - log(1-u) + logits) / temperature)
+        s = torch.log(u) - torch.log(1 - u) + self.qz_logits
+        s = torch.sigmoid(s / self.temperature)
+
+        # Stretch and clamp to create hard concrete
+        s = s * (self.zeta - self.gamma) + self.gamma
+        gates = torch.clamp(s, 0, 1)
+
+        return gates
+
+    def _deterministic_gates(self) -> torch.Tensor:
+        """
+        Compute deterministic gates for evaluation.
+
+        Returns
+        -------
+        torch.Tensor
+            Deterministic gate values in [0, 1]
+        """
+        # Use mean of the distribution
+        probs = torch.sigmoid(self.qz_logits)
+
+        # Apply stretching transformation
+        gates = probs * (self.zeta - self.gamma) + self.gamma
+
+        return torch.clamp(gates, 0, 1)
+
+    def get_penalty(self) -> torch.Tensor:
+        """
+        Compute the expected L0 norm (number of active gates).
+
+        Returns
+        -------
+        torch.Tensor
+            Expected number of non-zero gates
+        """
+        # Shift logits to account for hard concrete bounds
+        logits_shifted = self.qz_logits - self.temperature * math.log(
+            -self.gamma / self.zeta
+        )
+
+        # Probability that gate is active (non-zero)
+        prob_active = torch.sigmoid(logits_shifted)
+
+        return prob_active.sum()
+
+    def get_active_prob(self) -> torch.Tensor:
+        """
+        Get the probability that each gate is active.
+
+        Returns
+        -------
+        torch.Tensor
+            Probability of each gate being non-zero
+        """
+        logits_shifted = self.qz_logits - self.temperature * math.log(
+            -self.gamma / self.zeta
+        )
+        return torch.sigmoid(logits_shifted)
+
+    def get_sparsity(self) -> float:
+        """
+        Get the overall sparsity level (fraction of inactive gates).
+
+        Returns
+        -------
+        float
+            Sparsity level in [0, 1] where 1 means all gates are inactive
+        """
+        with torch.no_grad():
+            prob_active = self.get_active_prob()
+            return 1.0 - prob_active.mean().item()
+
+    @torch.no_grad()
+    def get_num_active(self) -> int:
+        """
+        Get the number of active gates (for monitoring).
+
+        Returns
+        -------
+        int
+            Number of gates that are likely active
+        """
+        prob_active = self.get_active_prob()
+        return int((prob_active > 0.5).sum().item())
+
+    def extra_repr(self) -> str:
+        """String representation for printing."""
+        return (
+            f"gate_size={self.gate_size}, temperature={self.temperature}, "
+            f"stretch={self.stretch}, init_mean={self.init_mean}"
+        )