l0-python 0.4.3__tar.gz → 0.6.0__tar.gz

{l0_python-0.4.3/l0_python.egg-info → l0_python-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: l0-python
-Version: 0.4.3
+Version: 0.6.0
 Summary: L0 regularization for sparse neural networks and intelligent sampling
 Author-email: PolicyEngine <hello@policyengine.org>
 License: MIT License

{l0_python-0.4.3 → l0_python-0.6.0}/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_python-0.4.3 → l0_python-0.6.0}/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()
@@ -304,6 +310,8 @@ class SparseCalibrationWeights(nn.Module):
         verbose: bool = False,
         verbose_freq: int = 100,
         target_groups: np.ndarray | None = None,
+        normalize_groups: bool = True,
+        group_multipliers: dict[int, float] | None = None,
     ) -> "SparseCalibrationWeights":
         """
         Fit calibration weights using gradient descent.
@@ -332,6 +340,12 @@ class SparseCalibrationWeights(nn.Module):
             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.
+        normalize_groups : bool
+            Whether to normalize within groups so each group contributes
+            equally to the loss. Default True (backward-compatible).
+        group_multipliers : dict[int, float], optional
+            Per-group loss scaling factors. Applied after normalization
+            (if enabled). Requires target_groups to be set.
 
         Returns
         -------
@@ -344,27 +358,53 @@ class SparseCalibrationWeights(nn.Module):
         # Convert M to torch sparse (will be cached)
         M_torch = self._convert_sparse_to_torch(M)
 
+        # Validate group_multipliers
+        if group_multipliers is not None and target_groups is None:
+            raise ValueError(
+                "group_multipliers requires target_groups to be set"
+            )
+
         # 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
+
+            if normalize_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()
+                    group_weights[group_mask] = 1.0 / group_size
+            else:
+                group_weights = torch.ones_like(y)
+
+            if group_multipliers is not None:
+                for gid, mult in group_multipliers.items():
+                    group_mask = target_groups == gid
+                    if not group_mask.any():
+                        raise ValueError(
+                            f"group_multipliers key {gid} not found "
+                            f"in target_groups"
+                        )
+                    group_weights[group_mask] *= mult
         else:
-            # No grouping - all targets weighted equally
             group_weights = torch.ones_like(y)
 
+        if verbose:
+            if target_groups is not None:
+                n_groups = len(torch.unique(target_groups))
+                parts = [f"{n_groups} groups"]
+                parts.append(
+                    f"normalize={'on' if normalize_groups else 'off'}"
+                )
+                if group_multipliers is not None:
+                    parts.append(f"multipliers={group_multipliers}")
+                print(f"Groups: {', '.join(parts)}")
+            else:
+                print("Groups: none (uniform weights)")
+
         # Add jitter to weights to break symmetry (if jitter_sd > 0)
         if self.log_weight_jitter_sd > 0:
             with torch.no_grad():
@@ -453,6 +493,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 +544,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 +553,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.3 → l0_python-0.6.0/l0_python.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: l0-python
-Version: 0.4.3
+Version: 0.6.0
 Summary: L0 regularization for sparse neural networks and intelligent sampling
 Author-email: PolicyEngine <hello@policyengine.org>
 License: MIT License

{l0_python-0.4.3 → l0_python-0.6.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "l0-python"
-version = "0.4.3"
+version = "0.6.0"
 description = "L0 regularization for sparse neural networks and intelligent sampling"
 readme = "README.md"
 authors = [

{l0_python-0.4.3 → l0_python-0.6.0}/tests/test_calibration.py

@@ -769,3 +769,224 @@ class TestSparseCalibrationWeights:
             log_alphas = model_extreme.log_alpha.numpy()
             # Should not have inf or -inf values
             assert np.all(np.isfinite(log_alphas))
+
+    def test_group_multipliers_without_normalization(self):
+        """Multipliers work with normalize_groups=False."""
+        N = 100
+        Q = 39  # 3 singletons + 18 + 18
+
+        np.random.seed(42)
+        torch.manual_seed(42)
+
+        M = sp.random(Q, N, density=0.3, format="csr")
+        y_singletons = np.array([1e9, 5e8, 2e9])
+        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])
+
+        target_groups = np.array([0, 1, 2] + [3] * 18 + [4] * 18)
+
+        # Without multipliers (baseline)
+        model_base = SparseCalibrationWeights(n_features=N)
+        model_base.fit(
+            M,
+            y,
+            lambda_l0=0.0001,
+            lr=0.1,
+            epochs=500,
+            loss_type="relative",
+            verbose=False,
+            target_groups=target_groups,
+            normalize_groups=False,
+        )
+
+        # With 10x multiplier on singletons
+        model_mult = SparseCalibrationWeights(n_features=N)
+        model_mult.fit(
+            M,
+            y,
+            lambda_l0=0.0001,
+            lr=0.1,
+            epochs=500,
+            loss_type="relative",
+            verbose=False,
+            target_groups=target_groups,
+            normalize_groups=False,
+            group_multipliers={0: 10.0, 1: 10.0, 2: 10.0},
+        )
+
+        with torch.no_grad():
+            y_pred_base = model_base.predict(M).cpu().numpy()
+            y_pred_mult = model_mult.predict(M).cpu().numpy()
+
+            # Singleton relative errors
+            sing_err_base = np.abs(
+                (y[:3] - y_pred_base[:3]) / (y[:3] + 1)
+            ).mean()
+            sing_err_mult = np.abs(
+                (y[:3] - y_pred_mult[:3]) / (y[:3] + 1)
+            ).mean()
+
+            assert sing_err_mult < sing_err_base * 1.1, (
+                f"10x multiplier should improve singleton fit: "
+                f"{sing_err_mult:.4f} vs {sing_err_base:.4f}"
+            )
+
+    def test_group_multipliers_with_normalization(self):
+        """Multipliers compose with normalization."""
+        N = 100
+        Q = 39
+
+        np.random.seed(42)
+        torch.manual_seed(42)
+
+        M = sp.random(Q, N, density=0.3, format="csr")
+        y_singletons = np.array([1e9, 5e8, 2e9])
+        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])
+
+        target_groups = np.array([0, 1, 2] + [3] * 18 + [4] * 18)
+
+        # With normalization only (baseline)
+        model_norm = SparseCalibrationWeights(n_features=N)
+        model_norm.fit(
+            M,
+            y,
+            lambda_l0=0.0001,
+            lr=0.1,
+            epochs=500,
+            loss_type="relative",
+            verbose=False,
+            target_groups=target_groups,
+            normalize_groups=True,
+        )
+
+        # With normalization + 5x multiplier on group 3
+        model_norm_mult = SparseCalibrationWeights(n_features=N)
+        model_norm_mult.fit(
+            M,
+            y,
+            lambda_l0=0.0001,
+            lr=0.1,
+            epochs=500,
+            loss_type="relative",
+            verbose=False,
+            target_groups=target_groups,
+            normalize_groups=True,
+            group_multipliers={3: 5.0},
+        )
+
+        with torch.no_grad():
+            y_pred_norm = model_norm.predict(M).cpu().numpy()
+            y_pred_mult = model_norm_mult.predict(M).cpu().numpy()
+
+            # Group 3 error should improve with multiplier
+            g3_err_norm = np.abs(
+                (y[3:21] - y_pred_norm[3:21]) / (y[3:21] + 1)
+            ).mean()
+            g3_err_mult = np.abs(
+                (y[3:21] - y_pred_mult[3:21]) / (y[3:21] + 1)
+            ).mean()
+
+            assert g3_err_mult < g3_err_norm * 1.1, (
+                f"5x multiplier should improve group 3 fit: "
+                f"{g3_err_mult:.4f} vs {g3_err_norm:.4f}"
+            )
+
+    def test_group_multipliers_requires_target_groups(self):
+        """ValueError when group_multipliers set but target_groups is None."""
+        N = 50
+        Q = 10
+
+        M = sp.random(Q, N, density=0.3, format="csr")
+        y = np.random.uniform(100, 1000, size=Q)
+
+        model = SparseCalibrationWeights(n_features=N)
+        with pytest.raises(
+            ValueError,
+            match="group_multipliers requires target_groups",
+        ):
+            model.fit(
+                M,
+                y,
+                epochs=10,
+                verbose=False,
+                group_multipliers={0: 2.0},
+            )
+
+    def test_group_multipliers_invalid_group_id(self):
+        """ValueError when multiplier key doesn't exist in target_groups."""
+        N = 50
+        Q = 10
+
+        M = sp.random(Q, N, density=0.3, format="csr")
+        y = np.random.uniform(100, 1000, size=Q)
+        target_groups = np.zeros(Q, dtype=int)  # All group 0
+
+        model = SparseCalibrationWeights(n_features=N)
+        with pytest.raises(
+            ValueError,
+            match="group_multipliers key 99 not found",
+        ):
+            model.fit(
+                M,
+                y,
+                epochs=10,
+                verbose=False,
+                target_groups=target_groups,
+                group_multipliers={99: 2.0},
+            )
+
+    def test_normalize_groups_false_uniform_weights(self):
+        """normalize_groups=False with target_groups gives uniform weights."""
+        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)
+
+        # Without target_groups at all
+        model_none = SparseCalibrationWeights(n_features=N)
+        model_none.fit(
+            M,
+            y,
+            lambda_l0=0.0001,
+            lr=0.1,
+            epochs=500,
+            loss_type="relative",
+            verbose=False,
+            target_groups=None,
+        )
+
+        # With target_groups but normalize_groups=False (should be equivalent)
+        torch.manual_seed(42)
+        target_groups = np.array([0] * 5 + [1] * 5)
+        model_no_norm = SparseCalibrationWeights(n_features=N)
+        model_no_norm.fit(
+            M,
+            y,
+            lambda_l0=0.0001,
+            lr=0.1,
+            epochs=500,
+            loss_type="relative",
+            verbose=False,
+            target_groups=target_groups,
+            normalize_groups=False,
+        )
+
+        with torch.no_grad():
+            y_pred_none = model_none.predict(M).cpu().numpy()
+            y_pred_no_norm = model_no_norm.predict(M).cpu().numpy()
+
+            err_none = np.abs((y - y_pred_none) / (y + 1)).mean()
+            err_no_norm = np.abs((y - y_pred_no_norm) / (y + 1)).mean()
+
+            # Both should achieve similar error levels
+            assert abs(err_none - err_no_norm) < 0.1, (
+                f"normalize_groups=False should behave like no groups: "
+                f"{err_none:.4f} vs {err_no_norm:.4f}"
+            )