l0-python 0.1.0__py3-none-any.whl

l0/penalties.py

@@ -0,0 +1,381 @@
+"""
+Penalty computation and management utilities for L0 regularization.
+
+Provides functions for computing penalties, tracking sparsity,
+and managing temperature schedules.
+"""
+
+from collections import defaultdict
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+def compute_l0_penalty(model: nn.Module) -> torch.Tensor:
+    """
+    Compute total L0 penalty across all L0 layers in a model.
+
+    Parameters
+    ----------
+    model : nn.Module
+        Model containing L0 layers
+
+    Returns
+    -------
+    torch.Tensor
+        Total L0 penalty (expected number of active parameters)
+    """
+    penalty = torch.tensor(0.0)
+
+    for module in model.modules():
+        if hasattr(module, "get_l0_penalty"):
+            penalty = penalty + module.get_l0_penalty()  # type: ignore[operator]
+
+    return penalty
+
+
+def compute_l2_penalty(model: nn.Module) -> torch.Tensor:
+    """
+    Compute L2 penalty (weight decay) across all layers.
+
+    Parameters
+    ----------
+    model : nn.Module
+        Model to compute L2 penalty for
+
+    Returns
+    -------
+    torch.Tensor
+        Total L2 penalty (sum of squared weights)
+    """
+    penalty = torch.tensor(0.0)
+
+    for name, param in model.named_parameters():
+        if "weight" in name and param.requires_grad:
+            penalty = penalty + (param**2).sum()
+
+    return penalty
+
+
+def compute_l0l2_penalty(
+    model: nn.Module,
+    l0_lambda: float = 1e-3,
+    l2_lambda: float = 1e-4,
+) -> torch.Tensor:
+    """
+    Compute combined L0L2 penalty.
+
+    This combination is recommended in practice to prevent overfitting
+    while maintaining sparsity.
+
+    Parameters
+    ----------
+    model : nn.Module
+        Model to compute penalties for
+    l0_lambda : float
+        Weight for L0 penalty
+    l2_lambda : float
+        Weight for L2 penalty
+
+    Returns
+    -------
+    torch.Tensor
+        Combined penalty value
+    """
+    l0_penalty = compute_l0_penalty(model)
+    l2_penalty = compute_l2_penalty(model)
+
+    return l0_lambda * l0_penalty + l2_lambda * l2_penalty
+
+
+def get_sparsity_stats(model: nn.Module) -> dict[str, dict[str, Any]]:
+    """
+    Get sparsity statistics for all L0 layers in a model.
+
+    Parameters
+    ----------
+    model : nn.Module
+        Model containing L0 layers
+
+    Returns
+    -------
+    Dict[str, Dict[str, Any]]
+        Dictionary mapping layer names to their statistics
+    """
+    stats = {}
+
+    for name, module in model.named_modules():
+        if hasattr(module, "get_sparsity") and hasattr(
+            module, "get_l0_penalty"
+        ):
+            active_params = module.get_l0_penalty().item()
+
+            # Calculate total parameters
+            if hasattr(module, "weight"):
+                total_params = module.weight.numel()
+            elif hasattr(module, "gates"):
+                total_params = module.gates.gate_size
+                if isinstance(total_params, tuple):
+                    total_params = np.prod(total_params)
+            else:
+                total_params = active_params / (1 - module.get_sparsity())
+
+            stats[name] = {
+                "sparsity": module.get_sparsity(),
+                "active_params": active_params,
+                "total_params": total_params,
+            }
+
+    return stats
+
+
+def get_active_parameter_count(model: nn.Module) -> int:
+    """
+    Get the total number of active parameters in the model.
+
+    Parameters
+    ----------
+    model : nn.Module
+        Model containing L0 layers
+
+    Returns
+    -------
+    int
+        Total number of active parameters
+    """
+    total = 0
+
+    for module in model.modules():
+        if hasattr(module, "get_l0_penalty"):
+            total += int(module.get_l0_penalty().item())  # type: ignore[operator]
+
+    return total
+
+
+class TemperatureScheduler:
+    """
+    Temperature scheduler for annealing during training.
+
+    Gradually decreases temperature to make gates more discrete.
+
+    Parameters
+    ----------
+    initial_temp : float
+        Starting temperature
+    final_temp : float
+        Final temperature
+    anneal_epochs : int
+        Number of epochs over which to anneal
+    schedule : str
+        Type of schedule ('linear', 'exponential', 'cosine')
+    """
+
+    def __init__(
+        self,
+        initial_temp: float = 2.0,
+        final_temp: float = 0.1,
+        anneal_epochs: int = 100,
+        schedule: str = "exponential",
+    ):
+        self.initial_temp = initial_temp
+        self.final_temp = final_temp
+        self.anneal_epochs = anneal_epochs
+        self.schedule = schedule
+
+    def get_temperature(self, epoch: int) -> float:
+        """
+        Get temperature for a given epoch.
+
+        Parameters
+        ----------
+        epoch : int
+            Current epoch number
+
+        Returns
+        -------
+        float
+            Temperature value
+        """
+        if epoch >= self.anneal_epochs:
+            return self.final_temp
+
+        progress = epoch / self.anneal_epochs
+
+        if self.schedule == "linear":
+            temp = (
+                self.initial_temp
+                - (self.initial_temp - self.final_temp) * progress
+            )
+
+        elif self.schedule == "exponential":
+            log_temp = (
+                np.log(self.initial_temp)
+                - (np.log(self.initial_temp) - np.log(self.final_temp))
+                * progress
+            )
+            temp = np.exp(log_temp)
+
+        elif self.schedule == "cosine":
+            temp = (
+                self.final_temp
+                + (self.initial_temp - self.final_temp)
+                * (1 + np.cos(np.pi * progress))
+                / 2
+            )
+
+        else:
+            raise ValueError(f"Unknown schedule type: {self.schedule}")
+
+        return float(temp)
+
+
+def update_temperatures(model: nn.Module, temperature: float) -> None:
+    """
+    Update temperature for all HardConcrete distributions in model.
+
+    Parameters
+    ----------
+    model : nn.Module
+        Model containing HardConcrete gates
+    temperature : float
+        New temperature value
+    """
+    for module in model.modules():
+        # Check for HardConcrete gates in layers
+        if hasattr(module, "weight_gates"):
+            module.weight_gates.temperature = temperature  # type: ignore[union-attr]
+        if hasattr(module, "channel_gates"):
+            module.channel_gates.temperature = temperature  # type: ignore[union-attr]
+        if hasattr(module, "gates"):
+            if hasattr(module.gates, "temperature"):
+                module.gates.temperature = temperature
+
+        # Direct HardConcrete modules
+        from .distributions import HardConcrete
+
+        if isinstance(module, HardConcrete):
+            module.temperature = temperature
+
+
+class PenaltyTracker:
+    """
+    Track penalties and sparsity metrics during training.
+
+    Useful for monitoring and visualization.
+    """
+
+    def __init__(self) -> None:
+        self.history: dict[str, list[float]] = defaultdict(list)
+
+    def log(self, name: str, value: float) -> None:
+        """
+        Log a metric value.
+
+        Parameters
+        ----------
+        name : str
+            Metric name
+        value : float
+            Metric value
+        """
+        self.history[name].append(value)
+
+    def get_history(self, name: str) -> list[float]:
+        """
+        Get history for a metric.
+
+        Parameters
+        ----------
+        name : str
+            Metric name
+
+        Returns
+        -------
+        List[float]
+            History of values
+        """
+        return self.history[name]
+
+    def get_stats(self, name: str) -> dict[str, float]:
+        """
+        Get statistics for a metric.
+
+        Parameters
+        ----------
+        name : str
+            Metric name
+
+        Returns
+        -------
+        Dict[str, float]
+            Statistics (mean, min, max, last)
+        """
+        values = self.history[name]
+        if not values:
+            return {}
+
+        return {
+            "mean": float(np.mean(values)),
+            "std": float(np.std(values)),
+            "min": float(np.min(values)),
+            "max": float(np.max(values)),
+            "last": values[-1],
+        }
+
+    def clear(self) -> None:
+        """Clear all history."""
+        self.history.clear()
+
+    def save_to_csv(self, filepath: str) -> None:
+        """
+        Save history to CSV file.
+
+        Parameters
+        ----------
+        filepath : str
+            Path to save CSV
+        """
+        import pandas as pd
+
+        df = pd.DataFrame(self.history)
+        df.to_csv(filepath, index=False)
+
+    def plot_history(
+        self,
+        metrics: list[str] | None = None,
+        save_path: str | None = None,
+    ) -> None:
+        """
+        Plot metric history.
+
+        Parameters
+        ----------
+        metrics : Optional[List[str]]
+            Metrics to plot (all if None)
+        save_path : Optional[str]
+            Path to save figure
+        """
+        import matplotlib.pyplot as plt
+
+        metrics = metrics or list(self.history.keys())
+
+        fig, axes = plt.subplots(
+            len(metrics), 1, figsize=(10, 4 * len(metrics)), squeeze=False
+        )
+
+        for i, metric in enumerate(metrics):
+            values = self.history[metric]
+            axes[i, 0].plot(values)
+            axes[i, 0].set_title(metric)
+            axes[i, 0].set_xlabel("Step")
+            axes[i, 0].set_ylabel("Value")
+            axes[i, 0].grid(True)
+
+        plt.tight_layout()
+
+        if save_path:
+            plt.savefig(save_path)
+        else:
+            plt.show()

l0_python-0.1.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,10 @@
+l0/__init__.py,sha256=meL6jJ1CxO6pO373nY-nwGXnpzWZGC1Ts6q2TLyAeLM,1122
+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_python-0.1.0.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
+l0_python-0.1.0.dist-info/METADATA,sha256=hrAugZ-F7oKmleyJ8WjVTykRixBmuTiQ2M4oae1Ceu4,6794
+l0_python-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+l0_python-0.1.0.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
+l0_python-0.1.0.dist-info/RECORD,,

l0_python-0.1.0.dist-info/WHEEL

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

l0_python-0.1.0.dist-info/licenses/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.dist-info/top_level.txt

@@ -0,0 +1 @@
+l0