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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: l0-python
-Version: 0.1.1
+Version: 0.3.1
 Summary: L0 regularization for sparse neural networks and intelligent sampling
 Author-email: PolicyEngine <hello@policyengine.org>
 License: MIT License
@@ -36,7 +36,7 @@ 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
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0.0
@@ -46,11 +46,10 @@ 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"
 Requires-Dist: scikit-learn>=1.0; extra == "dev"
+Requires-Dist: yaml-changelog>=0.3.0; extra == "dev"
+Requires-Dist: twine>=4.0.0; extra == "dev"
 Provides-Extra: docs
 Requires-Dist: mystmd>=1.3; extra == "docs"
 Requires-Dist: jupyter>=1.0; extra == "docs"

l0_python-0.3.1/l0/calibration.py

@@ -0,0 +1,495 @@
+"""
+L0-regularized calibration weights for survey data.
+
+Implements sparse calibration weights that are constrained to be non-negative,
+suitable for survey weighting and importance sampling applications.
+"""
+
+# ruff: noqa: N803 N806
+
+import torch
+import torch.nn as nn
+import numpy as np
+from scipy import sparse as sp
+
+
+class SparseCalibrationWeights(nn.Module):
+    """
+    L0-regularized calibration weights with positivity constraint.
+
+    Designed for survey calibration where weights must be non-negative
+    and sparsity is desired to select a subset of households/units.
+
+    Parameters
+    ----------
+    n_features : int
+        Number of calibration weights (e.g., number of households)
+    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 or array-like
+        Initial probability of keeping each weight active.
+        If float, all weights use the same probability.
+        If array, must have shape (n_features,)
+    init_weights : float, array-like, or None
+        Initial weight values (on original scale, not log).
+        If float, all weights initialized to this value.
+        If array, must have shape (n_features,).
+        If None, defaults to 1.0 for all weights.
+    log_weight_jitter_sd : float
+        Standard deviation of noise added to log weights at start of fit() to break symmetry.
+        Set to 0 to disable jitter. Default is 0.0 (no jitter).
+    log_alpha_jitter_sd : float
+        Standard deviation of noise added to log_alpha at initialization to break symmetry.
+        Set to 0 to disable jitter. Default is 0.01 (following Louizos et al.).
+    device : str or torch.device
+        Device to run computations on ('cpu' or 'cuda')
+    """
+
+    def __init__(
+        self,
+        n_features: int,
+        beta: float = 2 / 3,
+        gamma: float = -0.1,
+        zeta: float = 1.1,
+        init_keep_prob: float | np.ndarray = 0.5,
+        init_weights: float | np.ndarray | None = None,
+        log_weight_jitter_sd: float = 0.0,
+        log_alpha_jitter_sd: float = 0.01,
+        device: str | torch.device = "cpu",
+    ):
+        super().__init__()
+        self.n_features = n_features
+        self.beta = beta
+        self.gamma = gamma
+        self.zeta = zeta
+        self.log_weight_jitter_sd = log_weight_jitter_sd
+        self.log_alpha_jitter_sd = log_alpha_jitter_sd
+        self.device = torch.device(device)
+
+        # Initialize weights (on original scale)
+        if init_weights is None:
+            # Default: all weights start at 1.0
+            weight_values = torch.ones(n_features, device=self.device)
+        elif isinstance(init_weights, (int, float)):
+            # Scalar: all weights start at this value
+            weight_values = torch.full(
+                (n_features,), float(init_weights), device=self.device
+            )
+        else:
+            # Array: use provided values
+            weight_values = torch.tensor(
+                init_weights, dtype=torch.float32, device=self.device
+            )
+            if weight_values.shape != (n_features,):
+                raise ValueError(
+                    f"init_weights array must have shape ({n_features},), "
+                    f"got {weight_values.shape}"
+                )
+
+        # Convert to log space to ensure positivity via exp transformation
+        # Add small epsilon to avoid log(0)
+        self.log_weight = nn.Parameter(torch.log(weight_values + 1e-8))
+
+        # L0 gate parameters
+        # Handle init_keep_prob as scalar or array
+        if isinstance(init_keep_prob, (int, float)):
+            # Scalar: broadcast to all features
+            keep_prob_values = torch.full(
+                (n_features,), float(init_keep_prob), device=self.device
+            )
+        else:
+            # Array: use provided values
+            keep_prob_values = torch.tensor(
+                init_keep_prob, dtype=torch.float32, device=self.device
+            )
+            if keep_prob_values.shape != (n_features,):
+                raise ValueError(
+                    f"init_keep_prob array must have shape ({n_features},), "
+                    f"got {keep_prob_values.shape}"
+                )
+            # Clip to valid probability range to avoid log(0) or log(inf)
+            keep_prob_values = keep_prob_values.clamp(1e-6, 1 - 1e-6)
+
+        # Convert probabilities to log_alpha
+        mu = torch.log(keep_prob_values / (1 - keep_prob_values))
+        # Add jitter to break symmetry (if specified)
+        if self.log_alpha_jitter_sd > 0:
+            jitter = (
+                torch.randn(n_features, device=self.device)
+                * self.log_alpha_jitter_sd
+            )
+            self.log_alpha = nn.Parameter(mu + jitter)
+        else:
+            self.log_alpha = nn.Parameter(mu)
+
+        # Cache for sparse tensor conversion
+        self._cached_M_torch: torch.sparse.Tensor | None = None
+        self._cached_M_shape: tuple[int, int] | None = None
+
+    def _convert_sparse_to_torch(
+        self, M_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_M_torch is not None
+            and self._cached_M_shape == M_sparse.shape
+            and M_sparse.nnz == self._cached_M_torch._nnz()
+        ):
+            return self._cached_M_torch
+
+        M_coo = M_sparse.tocoo()
+        indices = torch.LongTensor(np.vstack([M_coo.row, M_coo.col])).to(
+            self.device
+        )
+        values = torch.FloatTensor(M_coo.data).to(self.device)
+        M_torch = torch.sparse_coo_tensor(
+            indices,
+            values,
+            M_sparse.shape,
+            dtype=torch.float32,
+            device=self.device,
+        )
+
+        # Cache for future use
+        self._cached_M_torch = M_torch
+        self._cached_M_shape = M_sparse.shape
+
+        return M_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)
+        s = (torch.log(u) - torch.log(1 - u) + self.log_alpha) / self.beta
+        s = torch.sigmoid(s)
+        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)."""
+        s = torch.sigmoid(self.log_alpha / self.beta)
+        s_bar = s * (self.zeta - self.gamma) + self.gamma
+        return s_bar.clamp(0, 1)
+
+    def get_weights(self, deterministic: bool = False) -> torch.Tensor:
+        """
+        Get effective calibration weights.
+
+        Parameters
+        ----------
+        deterministic : bool
+            Whether to use deterministic gates (for inference)
+
+        Returns
+        -------
+        torch.Tensor
+            Positive calibration weights with L0 sparsity applied
+        """
+        # Sample or get deterministic gates
+        if deterministic:
+            gates = self.get_deterministic_gates()
+        else:
+            gates = self._sample_gates()
+
+        # Apply exp transformation and gates together
+        # This ensures gradient flow matches successful implementations
+        return torch.exp(self.log_weight) * gates
+
+    def forward(
+        self,
+        M: sp.spmatrix,
+        deterministic: bool = False,
+    ) -> torch.Tensor:
+        """
+        Apply calibration weights to metric matrix.
+
+        Parameters
+        ----------
+        M : scipy.sparse matrix
+            Metric matrix of shape (n_targets, n_features)
+        deterministic : bool
+            Whether to use deterministic gates (for inference)
+
+        Returns
+        -------
+        torch.Tensor
+            Weighted sums of shape (n_targets,)
+        """
+        # Convert sparse matrix to torch
+        M_torch = self._convert_sparse_to_torch(M)
+
+        # Get effective weights
+        weights = self.get_weights(deterministic=deterministic)
+
+        # Sparse matrix multiplication
+        y = torch.sparse.mm(M_torch, weights.unsqueeze(1)).squeeze(1)
+
+        return y
+
+    def get_l0_penalty(self) -> torch.Tensor:
+        """
+        Compute L0 complexity penalty.
+
+        Returns expected number of active weights.
+        """
+        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_l2_penalty(self) -> torch.Tensor:
+        """
+        Compute L2 penalty on positive weights.
+
+        Helps prevent weight explosion when gates are partially open.
+
+        Returns
+        -------
+        torch.Tensor
+            L2 norm of positive weights
+        """
+        positive_weights = torch.exp(self.log_weight)
+        return (positive_weights**2).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).float().mean().item()
+
+    def get_active_weights(self) -> dict:
+        """
+        Get indices and values of active (non-zero) weights.
+
+        Returns
+        -------
+        dict
+            Dictionary with 'indices' and 'values' of active weights
+        """
+        with torch.no_grad():
+            weights = self.get_weights(deterministic=True)
+            active_mask = weights > 0
+
+            return {
+                "indices": torch.where(active_mask)[0],
+                "values": weights[active_mask],
+                "count": active_mask.sum().item(),
+            }
+
+    def fit(
+        self,
+        M: sp.spmatrix,
+        y: np.ndarray,
+        lambda_l0: float = 0.01,
+        lambda_l2: float = 0.0,
+        lr: float = 0.01,
+        epochs: int = 1000,
+        loss_type: str = "mse",
+        verbose: bool = False,
+        verbose_freq: int = 100,
+        target_groups: np.ndarray | None = None,
+    ) -> "SparseCalibrationWeights":
+        """
+        Fit calibration weights using gradient descent.
+
+        Parameters
+        ----------
+        M : scipy.sparse matrix
+            Metric matrix of shape (n_targets, n_features)
+        y : numpy.ndarray
+            Target values to match
+        lambda_l0 : float
+            L0 regularization strength
+        lambda_l2 : float
+            L2 regularization strength on positive weights
+        lr : float
+            Learning rate
+        epochs : int
+            Number of training epochs
+        loss_type : str
+            Type of loss function: 'mse' or 'relative'
+        verbose : bool
+            Whether to print progress
+        verbose_freq : int
+            How often to print progress
+        target_groups : numpy.ndarray, optional
+            Array of group IDs for each target. Targets in the same group
+            will be averaged together so each group contributes equally to loss.
+            If None, all targets are treated independently.
+
+        Returns
+        -------
+        self
+            Fitted model
+        """
+        # Convert y to tensor
+        y = torch.tensor(y, dtype=torch.float32, device=self.device)
+
+        # Convert M to torch sparse (will be cached)
+        M_torch = self._convert_sparse_to_torch(M)
+
+        # Compute group weights for loss averaging
+        if target_groups is not None:
+            # Convert to tensor
+            target_groups = torch.tensor(
+                target_groups, dtype=torch.long, device=self.device
+            )
+
+            # Calculate group weights: 1 / group_size for each target
+            unique_groups = torch.unique(target_groups)
+            group_weights = torch.zeros_like(y)
+
+            for group_id in unique_groups:
+                group_mask = target_groups == group_id
+                group_size = group_mask.sum().item()
+                # Each target in the group gets weight 1/group_size
+                # so the group's total contribution is 1
+                group_weights[group_mask] = 1.0 / group_size
+        else:
+            # No grouping - all targets weighted equally
+            group_weights = torch.ones_like(y)
+
+        # Add jitter to weights to break symmetry (if jitter_sd > 0)
+        if self.log_weight_jitter_sd > 0:
+            with torch.no_grad():
+                jitter = (
+                    torch.randn_like(self.log_weight)
+                    * self.log_weight_jitter_sd
+                )
+                self.log_weight.data += jitter
+
+        # Setup optimizer
+        optimizer = torch.optim.Adam([self.log_weight, self.log_alpha], lr=lr)
+
+        # Training loop
+        for epoch in range(epochs):
+            # Forward pass
+            y_pred = self.forward(M, deterministic=False)
+
+            # Compute loss with group weighting
+            if loss_type == "relative":
+                # Relative error: (y - y_pred)^2 / (y + 1)^2
+                # Adding 1 to avoid division by zero
+                relative_errors = (y - y_pred) / (y + 1)
+                # Apply group weights and then average
+                weighted_squared_errors = (
+                    relative_errors.pow(2) * group_weights
+                )
+                data_loss = (
+                    weighted_squared_errors.sum()
+                )  # Sum because weights already normalize
+            else:
+                # Standard MSE with group weighting
+                squared_errors = (y - y_pred).pow(2)
+                weighted_squared_errors = squared_errors * group_weights
+                data_loss = (
+                    weighted_squared_errors.sum()
+                )  # Sum because weights already normalize
+
+            l0_loss = self.get_l0_penalty()
+            loss = data_loss + lambda_l0 * l0_loss
+
+            # Add L2 penalty if specified
+            if lambda_l2 > 0:
+                l2_loss = self.get_l2_penalty()
+                loss = loss + lambda_l2 * l2_loss
+
+            # Backward pass
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            # Print progress
+            if verbose and (epoch + 1) % verbose_freq == 0:
+                with torch.no_grad():
+                    active_info = self.get_active_weights()
+                    weights = self.get_weights(deterministic=True)
+                    active_weights = weights[weights > 0]
+
+                    # Compute relative errors for meaningful output
+                    y_det = self.forward(M, deterministic=True)
+                    if loss_type == "relative":
+                        rel_errors = torch.abs((y - y_det) / (y + 1))
+                    else:
+                        # For MSE, show relative errors anyway for interpretability
+                        rel_errors = torch.abs((y - y_det) / (y + 1))
+
+                    # For reporting, we can show both overall and group-averaged errors
+                    mean_rel_err = rel_errors.mean().item()
+                    max_rel_err = rel_errors.max().item()
+
+                    # Compute mean group loss if groups are used
+                    if target_groups is not None:
+                        # Calculate mean loss per group
+                        group_losses = []
+                        for group_id in torch.unique(target_groups):
+                            group_mask = target_groups == group_id
+                            group_mean_err = (
+                                rel_errors[group_mask].mean().item()
+                            )
+                            group_losses.append(group_mean_err)
+                        mean_group_loss = np.mean(group_losses)
+                    else:
+                        mean_group_loss = mean_rel_err
+
+                    # Calculate sparsity percentage
+                    sparsity_pct = 100 * (
+                        1 - active_info["count"] / self.n_features
+                    )
+
+                    # Calculate components of the actual loss being minimized
+                    actual_data_loss = data_loss.item()
+                    actual_l0_loss = l0_loss.item()
+                    actual_total_loss = loss.item()
+
+                    if target_groups is not None:
+                        print(
+                            f"Epoch {epoch+1:4d}: "
+                            f"mean_group_loss={mean_group_loss:.1%}, "
+                            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)"
+                        )
+                    else:
+                        print(
+                            f"Epoch {epoch+1:4d}: "
+                            f"mean_error={mean_rel_err:.1%}, "
+                            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)"
+                        )
+
+        return self
+
+    def predict(self, M: sp.spmatrix) -> torch.Tensor:
+        """
+        Apply calibration weights to new data.
+
+        Parameters
+        ----------
+        M : scipy.sparse matrix
+            Metric matrix of shape (n_targets, n_features)
+
+        Returns
+        -------
+        torch.Tensor
+            Calibrated predictions
+        """
+        with torch.no_grad():
+            return self.forward(M, deterministic=True)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: l0-python
-Version: 0.1.1
+Version: 0.3.1
 Summary: L0 regularization for sparse neural networks and intelligent sampling
 Author-email: PolicyEngine <hello@policyengine.org>
 License: MIT License
@@ -36,7 +36,7 @@ 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
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0.0
@@ -46,11 +46,10 @@ 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"
 Requires-Dist: scikit-learn>=1.0; extra == "dev"
+Requires-Dist: yaml-changelog>=0.3.0; extra == "dev"
+Requires-Dist: twine>=4.0.0; extra == "dev"
 Provides-Extra: docs
 Requires-Dist: mystmd>=1.3; extra == "docs"
 Requires-Dist: jupyter>=1.0; extra == "docs"

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

@@ -2,6 +2,7 @@ LICENSE
 README.md
 pyproject.toml
 l0/__init__.py
+l0/calibration.py
 l0/distributions.py
 l0/gates.py
 l0/layers.py
@@ -12,6 +13,7 @@ l0_python.egg-info/SOURCES.txt
 l0_python.egg-info/dependency_links.txt
 l0_python.egg-info/requires.txt
 l0_python.egg-info/top_level.txt
+tests/test_calibration.py
 tests/test_distributions.py
 tests/test_gates.py
 tests/test_layers.py

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

@@ -6,11 +6,10 @@ scipy>=1.7.0
 pytest>=7.0
 pytest-cov>=3.0
 black>=22.0
-mypy>=0.900
-ruff>=0.0.280
-pre-commit>=3.0
 build>=1.3
 scikit-learn>=1.0
+yaml-changelog>=0.3.0
+twine>=4.0.0
 
 [docs]
 mystmd>=1.3

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

@@ -4,14 +4,14 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "l0-python"
-version = "0.1.1"
+version = "0.3.1"
 description = "L0 regularization for sparse neural networks and intelligent sampling"
 readme = "README.md"
 authors = [
     {name = "PolicyEngine", email = "hello@policyengine.org"},
 ]
 license = {file = "LICENSE"}
-requires-python = ">=3.13"
+requires-python = ">=3.11"
 classifiers = [
     "Development Status :: 4 - Beta",
     "Intended Audience :: Science/Research",
@@ -32,11 +32,10 @@ dev = [
     "pytest>=7.0",
     "pytest-cov>=3.0",
     "black>=22.0",
-    "mypy>=0.900",
-    "ruff>=0.0.280",
-    "pre-commit>=3.0",
     "build>=1.3",
     "scikit-learn>=1.0",
+    "yaml-changelog>=0.3.0",
+    "twine>=4.0.0",
 ]
 docs = [
     "mystmd>=1.3",
@@ -62,22 +61,6 @@ Issues = "https://github.com/PolicyEngine/L0/issues"
 line-length = 79
 target-version = ['py313']
 
-[tool.ruff]
-line-length = 79
-target-version = "py313"
-extend-exclude = ["docs/examples/*.ipynb"]
-
-[tool.ruff.lint]
-select = ["E", "F", "W", "I", "N", "UP"]
-ignore = ["E501", "N812", "UP038"]  # E501: Line length handled by black, N812: Allow F as alias for functional, UP038: Use tuple for isinstance
-
-[tool.mypy]
-python_version = "3.13"
-warn_return_any = true
-warn_unused_configs = true
-disallow_untyped_defs = true
-ignore_missing_imports = true
-
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 python_files = ["test_*.py"]

l0_python-0.3.1/tests/test_calibration.py

@@ -0,0 +1,545 @@
+"""
+Tests for SparseCalibrationWeights with positive weight constraints.
+"""
+
+import numpy as np
+import torch
+import pytest
+from scipy import sparse as sp
+from l0.calibration import SparseCalibrationWeights
+
+
+class TestSparseCalibrationWeights:
+    """Test suite for calibration weights with L0 sparsity."""
+
+    def test_positive_weights(self):
+        """Verify all weights remain non-negative."""
+        N = 100
+        Q = 20
+
+        # Create test data
+        M = sp.random(Q, N, density=0.3, format="csr")
+        y = np.random.randn(Q) + 10
+
+        model = SparseCalibrationWeights(n_features=N, init_weights=1.0)
+        model.fit(M, y, epochs=100, verbose=False)
+
+        # Check positivity
+        with torch.no_grad():
+            weights = model.get_weights(deterministic=True)
+            assert torch.all(weights >= 0), "Weights must be non-negative"
+
+    def test_sparse_ground_truth_relative_loss(self):
+        """Test recovery of sparse ground truth using relative loss."""
+        Q = 200  # targets
+        N = 2000  # features
+        N_active = 1000  # 50% sparsity
+
+        np.random.seed(42)
+        torch.manual_seed(42)
+
+        # Generate data with sparse ground truth
+        M_dense = np.random.lognormal(mean=1.5, sigma=0.25, size=(Q, N))
+        M = sp.csr_matrix(M_dense)
+
+        w_true = np.zeros(N)
+        active_indices = np.random.choice(N, size=N_active, replace=False)
+        w_true[active_indices] = np.random.lognormal(
+            mean=2.0, sigma=1.0, size=N_active
+        )
+
+        y = M @ w_true
+
+        # Fit with relative loss
+        model = SparseCalibrationWeights(
+            n_features=N,
+            beta=0.66,
+            gamma=-0.1,
+            zeta=1.1,
+            init_keep_prob=0.3,
+            init_weights=1.0,  # Start all weights at 1.0
+            log_weight_jitter_sd=0.5,  # Add jitter for symmetry breaking
+        )
+
+        model.fit(
+            M=M,
+            y=y,
+            lambda_l0=0.0005,  # Tuned for ~50% sparsity with relative loss
+            lambda_l2=1e-6,
+            lr=0.2,
+            epochs=2000,
+            loss_type="relative",
+            verbose=False,
+        )
+
+        # Check sparsity is reasonable (between 30% and 70%)
+        sparsity = model.get_sparsity()
+        assert (
+            0.3 <= sparsity <= 0.7
+        ), f"Sparsity {sparsity:.2%} not in expected range"
+
+        # Check relative loss is low
+        with torch.no_grad():
+            y_pred = model.predict(M).cpu().numpy()
+            rel_loss = np.mean(((y - y_pred) / (y + 1)) ** 2)
+            assert rel_loss < 0.1, f"Relative loss {rel_loss:.4f} too high"
+
+    def test_relative_vs_mse_loss(self):
+        """Compare relative loss vs MSE for large-scale data."""
+        Q = 100
+        N = 500
+
+        np.random.seed(123)
+        torch.manual_seed(123)
+
+        # Large-scale data
+        M = sp.random(Q, N, density=0.5, format="csr")
+        M.data = np.abs(M.data) * 1000  # Large values
+        y = np.random.uniform(1000, 100000, size=Q)
+
+        # Train with MSE
+        model_mse = SparseCalibrationWeights(n_features=N)
+        model_mse.fit(
+            M,
+            y,
+            lambda_l0=1e-10,  # Very small for MSE
+            lr=0.1,
+            epochs=500,
+            loss_type="mse",
+            verbose=False,
+        )
+
+        # Train with relative loss
+        model_rel = SparseCalibrationWeights(n_features=N)
+        model_rel.fit(
+            M,
+            y,
+            lambda_l0=0.001,  # Can use larger penalty
+            lr=0.1,
+            epochs=500,
+            loss_type="relative",
+            verbose=False,
+        )
+
+        # Relative loss should achieve better relative accuracy
+        with torch.no_grad():
+            y_pred_mse = model_mse.predict(M).cpu().numpy()
+            y_pred_rel = model_rel.predict(M).cpu().numpy()
+
+            rel_err_mse = np.mean(np.abs((y - y_pred_mse) / (y + 1)))
+            rel_err_rel = np.mean(np.abs((y - y_pred_rel) / (y + 1)))
+
+            # Relative loss should do better on relative error
+            assert (
+                rel_err_rel <= rel_err_mse * 1.5
+            ), f"Relative loss should handle scale better: {rel_err_rel:.4f} vs {rel_err_mse:.4f}"
+
+    def test_sparsity_control(self):
+        """Test that L0 penalty controls sparsity level."""
+        Q = 50
+        N = 200
+
+        np.random.seed(123)
+        torch.manual_seed(123)
+
+        M = sp.random(Q, N, density=0.3, format="csr")
+        y = np.random.randn(Q) + 10
+
+        sparsities = []
+
+        # Test different L0 penalties
+        for lambda_l0 in [0.0001, 0.001, 0.01]:
+            model = SparseCalibrationWeights(n_features=N, init_keep_prob=0.5)
+            model.fit(
+                M,
+                y,
+                lambda_l0=lambda_l0,
+                lr=0.1,
+                epochs=2000,
+                loss_type="relative",
+                verbose=False,
+            )
+            sparsities.append(model.get_sparsity())
+
+        # Higher penalty should give more sparsity
+        assert (
+            sparsities[0] < sparsities[1]
+        ), "Higher L0 penalty should increase sparsity"
+        assert (
+            sparsities[1] < sparsities[2]
+        ), "Higher L0 penalty should increase sparsity"
+
+    def test_get_active_weights(self):
+        """Test active weight extraction."""
+        N = 100
+        model = SparseCalibrationWeights(n_features=N, init_weights=1.0)
+
+        # Simple test data
+        M = sp.eye(N, format="csr")
+        y = np.ones(N)
+
+        model.fit(M, y, lambda_l0=0.01, epochs=100, verbose=False)
+
+        active_info = model.get_active_weights()
+
+        assert "indices" in active_info
+        assert "values" in active_info
+        assert "count" in active_info
+        assert active_info["count"] == len(active_info["indices"])
+        assert len(active_info["values"]) == active_info["count"]
+
+        # All active values should be positive
+        if active_info["count"] > 0:
+            assert torch.all(active_info["values"] > 0)
+
+    def test_deterministic_inference(self):
+        """Test that inference is deterministic."""
+        N = 50
+        Q = 10
+
+        np.random.seed(123)
+        torch.manual_seed(123)
+
+        M = sp.random(Q, N, density=0.5, format="csr")
+        y = np.random.randn(Q)
+
+        model = SparseCalibrationWeights(n_features=N, init_weights=1.0)
+        model.fit(M, y, epochs=100, verbose=False)
+
+        # Multiple predictions should be identical
+        with torch.no_grad():
+            pred1 = model.predict(M).cpu().numpy()
+            pred2 = model.predict(M).cpu().numpy()
+
+        np.testing.assert_array_equal(
+            pred1, pred2, "Predictions should be deterministic"
+        )
+
+    def test_l2_regularization(self):
+        """Test that L2 penalty prevents weight explosion."""
+        N = 100
+        Q = 20
+
+        np.random.seed(123)
+        torch.manual_seed(123)
+
+        M = sp.random(Q, N, density=0.3, format="csr")
+        y = np.random.randn(Q) * 100  # Large scale
+
+        # Train without L2
+        model_no_l2 = SparseCalibrationWeights(n_features=N)
+        model_no_l2.fit(
+            M, y, lambda_l0=0.0001, lambda_l2=0.0, epochs=200, verbose=False
+        )
+
+        # Train with L2
+        model_with_l2 = SparseCalibrationWeights(n_features=N)
+        model_with_l2.fit(
+            M, y, lambda_l0=0.0001, lambda_l2=0.01, epochs=200, verbose=False
+        )
+
+        with torch.no_grad():
+            weights_no_l2 = model_no_l2.get_weights(deterministic=True)
+            weights_with_l2 = model_with_l2.get_weights(deterministic=True)
+
+            # L2 should reduce weight magnitudes
+            assert (
+                weights_with_l2.max() <= weights_no_l2.max() * 2.0
+            ), "L2 should prevent extreme weights"
+
+    def test_group_wise_averaging(self):
+        """Test that group-wise averaging balances loss contributions."""
+        N = 100  # features (households)
+
+        # Create targets with different cardinalities:
+        # - 3 singleton targets (like national targets)
+        # - 18 targets in one group (like age bins for one state)
+        # - 18 targets in another group (like age bins for another state)
+        Q = 3 + 18 + 18  # 39 total targets
+
+        np.random.seed(42)
+        torch.manual_seed(42)
+
+        # Create matrix with varying scales
+        M = sp.random(Q, N, density=0.3, format="csr")
+
+        # Create target values with different scales
+        # Singletons: large values (billions scale)
+        y_singletons = np.array([1e9, 5e8, 2e9])
+        # Groups: smaller values (thousands scale)
+        y_group1 = np.random.uniform(1e3, 1e6, size=18)
+        y_group2 = np.random.uniform(1e3, 1e6, size=18)
+        y = np.concatenate([y_singletons, y_group1, y_group2])
+
+        # Create target groups
+        # Groups 0, 1, 2: singletons (each national target)
+        # Group 3: all 18 targets from first age group
+        # Group 4: all 18 targets from second age group
+        target_groups = np.array(
+            [0, 1, 2]  # 3 singletons
+            + [3] * 18  # Group 3
+            + [4] * 18  # Group 4
+        )
+
+        # Train WITHOUT grouping (baseline)
+        model_no_groups = SparseCalibrationWeights(n_features=N)
+        model_no_groups.fit(
+            M,
+            y,
+            lambda_l0=0.0001,
+            lr=0.1,
+            epochs=500,
+            loss_type="relative",
+            verbose=False,
+            target_groups=None,  # No grouping
+        )
+
+        # Train WITH grouping
+        model_with_groups = SparseCalibrationWeights(n_features=N)
+        model_with_groups.fit(
+            M,
+            y,
+            lambda_l0=0.0001,
+            lr=0.1,
+            epochs=500,
+            loss_type="relative",
+            verbose=False,
+            target_groups=target_groups,
+        )
+
+        # Compute errors by group
+        with torch.no_grad():
+            y_pred_no_groups = model_no_groups.predict(M).cpu().numpy()
+            y_pred_with_groups = model_with_groups.predict(M).cpu().numpy()
+
+            # Relative errors
+            rel_err_no_groups = np.abs((y - y_pred_no_groups) / (y + 1))
+            rel_err_with_groups = np.abs((y - y_pred_with_groups) / (y + 1))
+
+            # Average errors by group
+            singleton_err_no_groups = rel_err_no_groups[:3].mean()
+            group3_err_no_groups = rel_err_no_groups[3:21].mean()
+            group4_err_no_groups = rel_err_no_groups[21:].mean()
+
+            singleton_err_with_groups = rel_err_with_groups[:3].mean()
+            group3_err_with_groups = rel_err_with_groups[3:21].mean()
+            group4_err_with_groups = rel_err_with_groups[21:].mean()
+
+            # With grouping, singleton errors should be much better
+            # (they're not dominated by the 36 histogram targets)
+            assert singleton_err_with_groups < singleton_err_no_groups * 1.5, (
+                f"Grouping should improve singleton accuracy: "
+                f"{singleton_err_with_groups:.4f} vs {singleton_err_no_groups:.4f}"
+            )
+
+            # All groups should have relatively balanced errors with grouping
+            all_group_errors = [
+                singleton_err_with_groups,
+                group3_err_with_groups,
+                group4_err_with_groups,
+            ]
+            max_err = max(all_group_errors)
+            min_err = min(all_group_errors)
+
+            # Errors should be within an order of magnitude of each other
+            assert max_err < min_err * 10, (
+                f"Group errors should be balanced: "
+                f"min={min_err:.4f}, max={max_err:.4f}"
+            )
+
+    def test_group_wise_averaging_edge_cases(self):
+        """Test edge cases for group-wise averaging."""
+        N = 50
+        Q = 10
+
+        np.random.seed(42)
+        torch.manual_seed(42)
+
+        M = sp.random(Q, N, density=0.3, format="csr")
+        y = np.random.uniform(100, 1000, size=Q)
+
+        model = SparseCalibrationWeights(n_features=N, init_weights=1.0)
+
+        # Test 1: All targets in one group (should behave like no grouping)
+        target_groups_single = np.zeros(Q, dtype=int)
+        model.fit(
+            M,
+            y,
+            lambda_l0=0.00001,  # Lower penalty for better convergence
+            epochs=2000,  # Plenty of epochs
+            lr=0.2,  # Higher learning rate
+            loss_type="relative",
+            verbose=False,
+            target_groups=target_groups_single,
+        )
+
+        with torch.no_grad():
+            y_pred = model.predict(M).cpu().numpy()
+            rel_err = np.mean(np.abs((y - y_pred) / (y + 1)))
+            assert (
+                rel_err < 0.5
+            ), f"Single group should still converge, got {rel_err:.4f}"
+
+        # Test 2: Each target in its own group (like all singletons)
+        target_groups_all_singleton = np.arange(Q)
+        model_new = SparseCalibrationWeights(n_features=N)
+        model_new.fit(
+            M,
+            y,
+            lambda_l0=0.00001,
+            epochs=2000,
+            lr=0.2,
+            loss_type="relative",
+            verbose=False,
+            target_groups=target_groups_all_singleton,
+        )
+
+        with torch.no_grad():
+            y_pred = model_new.predict(M).cpu().numpy()
+            rel_err = np.mean(np.abs((y - y_pred) / (y + 1)))
+            assert (
+                rel_err < 0.5
+            ), f"All singleton groups should converge, got {rel_err:.4f}"
+
+        # Test 3: Unbalanced groups (1 huge group, several small)
+        target_groups_unbalanced = np.array([0] * 7 + [1, 2, 3])
+        model_unbalanced = SparseCalibrationWeights(n_features=N)
+        model_unbalanced.fit(
+            M,
+            y,
+            lambda_l0=0.00001,
+            epochs=2000,
+            lr=0.2,
+            loss_type="relative",
+            verbose=False,
+            target_groups=target_groups_unbalanced,
+        )
+
+        with torch.no_grad():
+            y_pred = model_unbalanced.predict(M).cpu().numpy()
+            # Check that small groups aren't ignored
+            small_group_errors = np.abs((y[7:] - y_pred[7:]) / (y[7:] + 1))
+            assert (
+                np.mean(small_group_errors) < 0.5
+            ), "Small groups should not be ignored"
+
+    def test_init_weights_options(self):
+        """Test different weight initialization options."""
+        N = 50
+
+        # Test 1: Default (None) should give all weights = 1.0
+        model_default = SparseCalibrationWeights(n_features=N)
+        with torch.no_grad():
+            weights = torch.exp(model_default.log_weight)
+            assert torch.allclose(weights, torch.ones(N), atol=1e-6)
+
+        # Test 2: Scalar initialization
+        model_scalar = SparseCalibrationWeights(n_features=N, init_weights=2.5)
+        with torch.no_grad():
+            weights = torch.exp(model_scalar.log_weight)
+            assert torch.allclose(weights, torch.full((N,), 2.5), atol=1e-6)
+
+        # Test 3: Array initialization
+        init_array = np.random.uniform(0.5, 2.0, size=N)
+        model_array = SparseCalibrationWeights(
+            n_features=N, init_weights=init_array
+        )
+        with torch.no_grad():
+            weights = torch.exp(model_array.log_weight).cpu().numpy()
+            np.testing.assert_allclose(weights, init_array, rtol=1e-5)
+
+        # Test 4: Wrong shape should raise error
+        with pytest.raises(ValueError, match="must have shape"):
+            SparseCalibrationWeights(n_features=N, init_weights=np.ones(N + 1))
+
+    def test_weight_jitter(self):
+        """Test that weight jitter works correctly."""
+        N = 100
+        Q = 20
+
+        np.random.seed(42)
+        torch.manual_seed(42)
+
+        M = sp.random(Q, N, density=0.3, format="csr")
+        y = np.random.randn(Q) + 10
+
+        # Model with jitter
+        model_with_jitter = SparseCalibrationWeights(
+            n_features=N, init_weights=1.0, log_weight_jitter_sd=0.5
+        )
+
+        # Store initial weights
+        initial_weights = model_with_jitter.log_weight.data.clone()
+
+        # Fit should add jitter
+        model_with_jitter.fit(M, y, epochs=10, verbose=False)
+
+        # Weights should have changed due to jitter (and training)
+        final_weights = model_with_jitter.log_weight.data
+        assert not torch.allclose(initial_weights, final_weights)
+
+        # Model without jitter
+        torch.manual_seed(42)  # Reset seed
+        model_no_jitter = SparseCalibrationWeights(
+            n_features=N,
+            init_weights=1.0,
+            log_weight_jitter_sd=0.0,  # No jitter
+        )
+
+        initial_weights_no_jitter = model_no_jitter.log_weight.data.clone()
+        model_no_jitter.fit(M, y, epochs=1, verbose=False)  # Just 1 epoch
+
+        # After 1 epoch, change should be small without jitter
+        weights_after_1_epoch = model_no_jitter.log_weight.data
+        # The change is due to gradient updates only
+        change = (
+            (weights_after_1_epoch - initial_weights_no_jitter).abs().max()
+        )
+        assert change < 1.0, "Without jitter, initial change should be small"
+
+    def test_init_keep_prob_options(self):
+        """Test init_keep_prob as scalar and array."""
+        n_features = 20
+
+        # Test 1: Scalar init_keep_prob (existing behavior)
+        model_scalar = SparseCalibrationWeights(
+            n_features=n_features,
+            init_keep_prob=0.7,
+        )
+        # All log_alpha values should be similar (around log(0.7/0.3) plus small jitter)
+        expected_mu = np.log(0.7 / 0.3)
+        with torch.no_grad():
+            log_alphas = model_scalar.log_alpha.numpy()
+            # Check they're all close to expected value (within jitter range)
+            assert np.all(np.abs(log_alphas - expected_mu) < 0.1)
+
+        # Test 2: Array init_keep_prob
+        keep_probs = np.linspace(0.1, 0.9, n_features)
+        model_array = SparseCalibrationWeights(
+            n_features=n_features,
+            init_keep_prob=keep_probs,
+        )
+        # Each log_alpha should correspond to its keep_prob
+        with torch.no_grad():
+            log_alphas = model_array.log_alpha.numpy()
+            expected_mus = np.log(keep_probs / (1 - keep_probs))
+            # Check each is close to its expected value
+            assert np.all(np.abs(log_alphas - expected_mus) < 0.1)
+
+        # Test 3: Wrong shape should raise error
+        with pytest.raises(ValueError, match="must have shape"):
+            SparseCalibrationWeights(
+                n_features=10,
+                init_keep_prob=np.ones(5),  # Wrong size
+            )
+
+        # Test 4: Edge case probabilities get clamped
+        extreme_probs = np.array([0.0, 0.5, 1.0])
+        model_extreme = SparseCalibrationWeights(
+            n_features=3,
+            init_keep_prob=extreme_probs,
+        )
+        with torch.no_grad():
+            log_alphas = model_extreme.log_alpha.numpy()
+            # Should not have inf or -inf values
+            assert np.all(np.isfinite(log_alphas))