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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: l0-python
3
- Version: 0.1.1
3
+ Version: 0.3.0
4
4
  Summary: L0 regularization for sparse neural networks and intelligent sampling
5
5
  Author-email: PolicyEngine <hello@policyengine.org>
6
6
  License: MIT License
@@ -46,11 +46,10 @@ Provides-Extra: dev
46
46
  Requires-Dist: pytest>=7.0; extra == "dev"
47
47
  Requires-Dist: pytest-cov>=3.0; extra == "dev"
48
48
  Requires-Dist: black>=22.0; extra == "dev"
49
- Requires-Dist: mypy>=0.900; extra == "dev"
50
- Requires-Dist: ruff>=0.0.280; extra == "dev"
51
- Requires-Dist: pre-commit>=3.0; extra == "dev"
52
49
  Requires-Dist: build>=1.3; extra == "dev"
53
50
  Requires-Dist: scikit-learn>=1.0; extra == "dev"
51
+ Requires-Dist: yaml-changelog>=0.3.0; extra == "dev"
52
+ Requires-Dist: twine>=4.0.0; extra == "dev"
54
53
  Provides-Extra: docs
55
54
  Requires-Dist: mystmd>=1.3; extra == "docs"
56
55
  Requires-Dist: jupyter>=1.0; extra == "docs"
@@ -0,0 +1,436 @@
1
+ """
2
+ L0-regularized calibration weights for survey data.
3
+
4
+ Implements sparse calibration weights that are constrained to be non-negative,
5
+ suitable for survey weighting and importance sampling applications.
6
+ """
7
+
8
+ # ruff: noqa: N803 N806
9
+
10
+ import torch
11
+ import torch.nn as nn
12
+ import numpy as np
13
+ from scipy import sparse as sp
14
+
15
+
16
+ class SparseCalibrationWeights(nn.Module):
17
+ """
18
+ L0-regularized calibration weights with positivity constraint.
19
+
20
+ Designed for survey calibration where weights must be non-negative
21
+ and sparsity is desired to select a subset of households/units.
22
+
23
+ Parameters
24
+ ----------
25
+ n_features : int
26
+ Number of calibration weights (e.g., number of households)
27
+ beta : float
28
+ Temperature parameter for Hard Concrete distribution
29
+ gamma : float
30
+ Lower bound of stretched concrete distribution
31
+ zeta : float
32
+ Upper bound of stretched concrete distribution
33
+ init_keep_prob : float
34
+ Initial probability of keeping each weight active
35
+ init_weight_scale : float
36
+ Initial scale for log weights (controls initial weight magnitude)
37
+ device : str or torch.device
38
+ Device to run computations on ('cpu' or 'cuda')
39
+ """
40
+
41
+ def __init__(
42
+ self,
43
+ n_features: int,
44
+ beta: float = 2 / 3,
45
+ gamma: float = -0.1,
46
+ zeta: float = 1.1,
47
+ init_keep_prob: float = 0.5,
48
+ init_weight_scale: float = 1.0,
49
+ device: str | torch.device = "cpu",
50
+ ):
51
+ super().__init__()
52
+ self.n_features = n_features
53
+ self.beta = beta
54
+ self.gamma = gamma
55
+ self.zeta = zeta
56
+ self.device = torch.device(device)
57
+
58
+ # Log weights to ensure positivity via exp transformation
59
+ self.log_weight = nn.Parameter(
60
+ torch.normal(
61
+ mean=0.0,
62
+ std=init_weight_scale,
63
+ size=(n_features,),
64
+ device=self.device,
65
+ )
66
+ )
67
+
68
+ # L0 gate parameters
69
+ mu = torch.log(torch.tensor(init_keep_prob / (1 - init_keep_prob)))
70
+ self.log_alpha = nn.Parameter(
71
+ torch.normal(
72
+ mu.item(), 0.01, size=(n_features,), device=self.device
73
+ )
74
+ )
75
+
76
+ # Cache for sparse tensor conversion
77
+ self._cached_M_torch: torch.sparse.Tensor | None = None
78
+ self._cached_M_shape: tuple[int, int] | None = None
79
+
80
+ def _convert_sparse_to_torch(
81
+ self, M_sparse: sp.spmatrix
82
+ ) -> torch.sparse.Tensor:
83
+ """
84
+ Convert scipy sparse matrix to torch sparse tensor.
85
+
86
+ Caches the result if the shape matches to avoid redundant conversions.
87
+ """
88
+ if (
89
+ self._cached_M_torch is not None
90
+ and self._cached_M_shape == M_sparse.shape
91
+ and M_sparse.nnz == self._cached_M_torch._nnz()
92
+ ):
93
+ return self._cached_M_torch
94
+
95
+ M_coo = M_sparse.tocoo()
96
+ indices = torch.LongTensor(np.vstack([M_coo.row, M_coo.col])).to(
97
+ self.device
98
+ )
99
+ values = torch.FloatTensor(M_coo.data).to(self.device)
100
+ M_torch = torch.sparse_coo_tensor(
101
+ indices,
102
+ values,
103
+ M_sparse.shape,
104
+ dtype=torch.float32,
105
+ device=self.device,
106
+ )
107
+
108
+ # Cache for future use
109
+ self._cached_M_torch = M_torch
110
+ self._cached_M_shape = M_sparse.shape
111
+
112
+ return M_torch
113
+
114
+ def _sample_gates(self) -> torch.Tensor:
115
+ """Sample gates using Hard Concrete distribution."""
116
+ eps = 1e-6
117
+ u = torch.rand_like(self.log_alpha).clamp(eps, 1 - eps)
118
+ s = (torch.log(u) - torch.log(1 - u) + self.log_alpha) / self.beta
119
+ s = torch.sigmoid(s)
120
+ s_bar = s * (self.zeta - self.gamma) + self.gamma
121
+ return s_bar.clamp(0, 1)
122
+
123
+ def get_deterministic_gates(self) -> torch.Tensor:
124
+ """Get deterministic gate values (for inference)."""
125
+ s = torch.sigmoid(self.log_alpha / self.beta)
126
+ s_bar = s * (self.zeta - self.gamma) + self.gamma
127
+ return s_bar.clamp(0, 1)
128
+
129
+ def get_weights(self, deterministic: bool = False) -> torch.Tensor:
130
+ """
131
+ Get effective calibration weights.
132
+
133
+ Parameters
134
+ ----------
135
+ deterministic : bool
136
+ Whether to use deterministic gates (for inference)
137
+
138
+ Returns
139
+ -------
140
+ torch.Tensor
141
+ Positive calibration weights with L0 sparsity applied
142
+ """
143
+ # Sample or get deterministic gates
144
+ if deterministic:
145
+ gates = self.get_deterministic_gates()
146
+ else:
147
+ gates = self._sample_gates()
148
+
149
+ # Apply exp transformation and gates together
150
+ # This ensures gradient flow matches successful implementations
151
+ return torch.exp(self.log_weight) * gates
152
+
153
+ def forward(
154
+ self,
155
+ M: sp.spmatrix,
156
+ deterministic: bool = False,
157
+ ) -> torch.Tensor:
158
+ """
159
+ Apply calibration weights to metric matrix.
160
+
161
+ Parameters
162
+ ----------
163
+ M : scipy.sparse matrix
164
+ Metric matrix of shape (n_targets, n_features)
165
+ deterministic : bool
166
+ Whether to use deterministic gates (for inference)
167
+
168
+ Returns
169
+ -------
170
+ torch.Tensor
171
+ Weighted sums of shape (n_targets,)
172
+ """
173
+ # Convert sparse matrix to torch
174
+ M_torch = self._convert_sparse_to_torch(M)
175
+
176
+ # Get effective weights
177
+ weights = self.get_weights(deterministic=deterministic)
178
+
179
+ # Sparse matrix multiplication
180
+ y = torch.sparse.mm(M_torch, weights.unsqueeze(1)).squeeze(1)
181
+
182
+ return y
183
+
184
+ def get_l0_penalty(self) -> torch.Tensor:
185
+ """
186
+ Compute L0 complexity penalty.
187
+
188
+ Returns expected number of active weights.
189
+ """
190
+ c = -self.beta * torch.log(
191
+ torch.tensor(-self.gamma / self.zeta, device=self.device)
192
+ )
193
+ pi = torch.sigmoid(self.log_alpha + c)
194
+ return pi.sum()
195
+
196
+ def get_l2_penalty(self) -> torch.Tensor:
197
+ """
198
+ Compute L2 penalty on positive weights.
199
+
200
+ Helps prevent weight explosion when gates are partially open.
201
+
202
+ Returns
203
+ -------
204
+ torch.Tensor
205
+ L2 norm of positive weights
206
+ """
207
+ positive_weights = torch.exp(self.log_weight)
208
+ return (positive_weights**2).sum()
209
+
210
+ def get_sparsity(self) -> float:
211
+ """
212
+ Get current sparsity level.
213
+
214
+ Returns
215
+ -------
216
+ float
217
+ Fraction of weights that are effectively zero
218
+ """
219
+ with torch.no_grad():
220
+ gates = self.get_deterministic_gates()
221
+ return (gates == 0).float().mean().item()
222
+
223
+ def get_active_weights(self) -> dict:
224
+ """
225
+ Get indices and values of active (non-zero) weights.
226
+
227
+ Returns
228
+ -------
229
+ dict
230
+ Dictionary with 'indices' and 'values' of active weights
231
+ """
232
+ with torch.no_grad():
233
+ weights = self.get_weights(deterministic=True)
234
+ active_mask = weights > 0
235
+
236
+ return {
237
+ "indices": torch.where(active_mask)[0],
238
+ "values": weights[active_mask],
239
+ "count": active_mask.sum().item(),
240
+ }
241
+
242
+ def fit(
243
+ self,
244
+ M: sp.spmatrix,
245
+ y: np.ndarray,
246
+ lambda_l0: float = 0.01,
247
+ lambda_l2: float = 0.0,
248
+ lr: float = 0.01,
249
+ epochs: int = 1000,
250
+ loss_type: str = "mse",
251
+ verbose: bool = False,
252
+ verbose_freq: int = 100,
253
+ target_groups: np.ndarray | None = None,
254
+ ) -> "SparseCalibrationWeights":
255
+ """
256
+ Fit calibration weights using gradient descent.
257
+
258
+ Parameters
259
+ ----------
260
+ M : scipy.sparse matrix
261
+ Metric matrix of shape (n_targets, n_features)
262
+ y : numpy.ndarray
263
+ Target values to match
264
+ lambda_l0 : float
265
+ L0 regularization strength
266
+ lambda_l2 : float
267
+ L2 regularization strength on positive weights
268
+ lr : float
269
+ Learning rate
270
+ epochs : int
271
+ Number of training epochs
272
+ loss_type : str
273
+ Type of loss function: 'mse' or 'relative'
274
+ verbose : bool
275
+ Whether to print progress
276
+ verbose_freq : int
277
+ How often to print progress
278
+ target_groups : numpy.ndarray, optional
279
+ Array of group IDs for each target. Targets in the same group
280
+ will be averaged together so each group contributes equally to loss.
281
+ If None, all targets are treated independently.
282
+
283
+ Returns
284
+ -------
285
+ self
286
+ Fitted model
287
+ """
288
+ # Convert y to tensor
289
+ y = torch.tensor(y, dtype=torch.float32, device=self.device)
290
+
291
+ # Convert M to torch sparse (will be cached)
292
+ M_torch = self._convert_sparse_to_torch(M)
293
+
294
+ # Compute group weights for loss averaging
295
+ if target_groups is not None:
296
+ # Convert to tensor
297
+ target_groups = torch.tensor(
298
+ target_groups, dtype=torch.long, device=self.device
299
+ )
300
+
301
+ # Calculate group weights: 1 / group_size for each target
302
+ unique_groups = torch.unique(target_groups)
303
+ group_weights = torch.zeros_like(y)
304
+
305
+ for group_id in unique_groups:
306
+ group_mask = target_groups == group_id
307
+ group_size = group_mask.sum().item()
308
+ # Each target in the group gets weight 1/group_size
309
+ # so the group's total contribution is 1
310
+ group_weights[group_mask] = 1.0 / group_size
311
+ else:
312
+ # No grouping - all targets weighted equally
313
+ group_weights = torch.ones_like(y)
314
+
315
+ # Initialize weights
316
+ nn.init.normal_(self.log_weight, 0, 0.5)
317
+
318
+ # Setup optimizer
319
+ optimizer = torch.optim.Adam([self.log_weight, self.log_alpha], lr=lr)
320
+
321
+ # Training loop
322
+ for epoch in range(epochs):
323
+ # Forward pass
324
+ y_pred = self.forward(M, deterministic=False)
325
+
326
+ # Compute loss with group weighting
327
+ if loss_type == "relative":
328
+ # Relative error: (y - y_pred)^2 / (y + 1)^2
329
+ # Adding 1 to avoid division by zero
330
+ relative_errors = (y - y_pred) / (y + 1)
331
+ # Apply group weights and then average
332
+ weighted_squared_errors = (
333
+ relative_errors.pow(2) * group_weights
334
+ )
335
+ data_loss = (
336
+ weighted_squared_errors.sum()
337
+ ) # Sum because weights already normalize
338
+ else:
339
+ # Standard MSE with group weighting
340
+ squared_errors = (y - y_pred).pow(2)
341
+ weighted_squared_errors = squared_errors * group_weights
342
+ data_loss = (
343
+ weighted_squared_errors.sum()
344
+ ) # Sum because weights already normalize
345
+
346
+ l0_loss = self.get_l0_penalty()
347
+ loss = data_loss + lambda_l0 * l0_loss
348
+
349
+ # Add L2 penalty if specified
350
+ if lambda_l2 > 0:
351
+ l2_loss = self.get_l2_penalty()
352
+ loss = loss + lambda_l2 * l2_loss
353
+
354
+ # Backward pass
355
+ optimizer.zero_grad()
356
+ loss.backward()
357
+ optimizer.step()
358
+
359
+ # Print progress
360
+ if verbose and (epoch + 1) % verbose_freq == 0:
361
+ with torch.no_grad():
362
+ active_info = self.get_active_weights()
363
+ weights = self.get_weights(deterministic=True)
364
+ active_weights = weights[weights > 0]
365
+
366
+ # Compute relative errors for meaningful output
367
+ y_det = self.forward(M, deterministic=True)
368
+ if loss_type == "relative":
369
+ rel_errors = torch.abs((y - y_det) / (y + 1))
370
+ else:
371
+ # For MSE, show relative errors anyway for interpretability
372
+ rel_errors = torch.abs((y - y_det) / (y + 1))
373
+
374
+ # For reporting, we can show both overall and group-averaged errors
375
+ mean_rel_err = rel_errors.mean().item()
376
+ max_rel_err = rel_errors.max().item()
377
+
378
+ # Compute mean group loss if groups are used
379
+ if target_groups is not None:
380
+ # Calculate mean loss per group
381
+ group_losses = []
382
+ for group_id in torch.unique(target_groups):
383
+ group_mask = target_groups == group_id
384
+ group_mean_err = (
385
+ rel_errors[group_mask].mean().item()
386
+ )
387
+ group_losses.append(group_mean_err)
388
+ mean_group_loss = np.mean(group_losses)
389
+ else:
390
+ mean_group_loss = mean_rel_err
391
+
392
+ # Calculate sparsity percentage
393
+ sparsity_pct = 100 * (
394
+ 1 - active_info["count"] / self.n_features
395
+ )
396
+
397
+ # Calculate components of the actual loss being minimized
398
+ actual_data_loss = data_loss.item()
399
+ actual_l0_loss = l0_loss.item()
400
+ actual_total_loss = loss.item()
401
+
402
+ if target_groups is not None:
403
+ print(
404
+ f"Epoch {epoch+1:4d}: "
405
+ f"mean_group_loss={mean_group_loss:.1%}, "
406
+ f"max_error={max_rel_err:.1%}, "
407
+ f"total_loss={actual_total_loss:.3f}, "
408
+ f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)"
409
+ )
410
+ else:
411
+ print(
412
+ f"Epoch {epoch+1:4d}: "
413
+ f"mean_error={mean_rel_err:.1%}, "
414
+ f"max_error={max_rel_err:.1%}, "
415
+ f"total_loss={actual_total_loss:.3f}, "
416
+ f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)"
417
+ )
418
+
419
+ return self
420
+
421
+ def predict(self, M: sp.spmatrix) -> torch.Tensor:
422
+ """
423
+ Apply calibration weights to new data.
424
+
425
+ Parameters
426
+ ----------
427
+ M : scipy.sparse matrix
428
+ Metric matrix of shape (n_targets, n_features)
429
+
430
+ Returns
431
+ -------
432
+ torch.Tensor
433
+ Calibrated predictions
434
+ """
435
+ with torch.no_grad():
436
+ return self.forward(M, deterministic=True)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: l0-python
3
- Version: 0.1.1
3
+ Version: 0.3.0
4
4
  Summary: L0 regularization for sparse neural networks and intelligent sampling
5
5
  Author-email: PolicyEngine <hello@policyengine.org>
6
6
  License: MIT License
@@ -46,11 +46,10 @@ Provides-Extra: dev
46
46
  Requires-Dist: pytest>=7.0; extra == "dev"
47
47
  Requires-Dist: pytest-cov>=3.0; extra == "dev"
48
48
  Requires-Dist: black>=22.0; extra == "dev"
49
- Requires-Dist: mypy>=0.900; extra == "dev"
50
- Requires-Dist: ruff>=0.0.280; extra == "dev"
51
- Requires-Dist: pre-commit>=3.0; extra == "dev"
52
49
  Requires-Dist: build>=1.3; extra == "dev"
53
50
  Requires-Dist: scikit-learn>=1.0; extra == "dev"
51
+ Requires-Dist: yaml-changelog>=0.3.0; extra == "dev"
52
+ Requires-Dist: twine>=4.0.0; extra == "dev"
54
53
  Provides-Extra: docs
55
54
  Requires-Dist: mystmd>=1.3; extra == "docs"
56
55
  Requires-Dist: jupyter>=1.0; extra == "docs"
@@ -2,6 +2,7 @@ LICENSE
2
2
  README.md
3
3
  pyproject.toml
4
4
  l0/__init__.py
5
+ l0/calibration.py
5
6
  l0/distributions.py
6
7
  l0/gates.py
7
8
  l0/layers.py
@@ -12,6 +13,7 @@ l0_python.egg-info/SOURCES.txt
12
13
  l0_python.egg-info/dependency_links.txt
13
14
  l0_python.egg-info/requires.txt
14
15
  l0_python.egg-info/top_level.txt
16
+ tests/test_calibration.py
15
17
  tests/test_distributions.py
16
18
  tests/test_gates.py
17
19
  tests/test_layers.py
@@ -6,11 +6,10 @@ scipy>=1.7.0
6
6
  pytest>=7.0
7
7
  pytest-cov>=3.0
8
8
  black>=22.0
9
- mypy>=0.900
10
- ruff>=0.0.280
11
- pre-commit>=3.0
12
9
  build>=1.3
13
10
  scikit-learn>=1.0
11
+ yaml-changelog>=0.3.0
12
+ twine>=4.0.0
14
13
 
15
14
  [docs]
16
15
  mystmd>=1.3
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "l0-python"
7
- version = "0.1.1"
7
+ version = "0.3.0"
8
8
  description = "L0 regularization for sparse neural networks and intelligent sampling"
9
9
  readme = "README.md"
10
10
  authors = [
@@ -32,11 +32,10 @@ dev = [
32
32
  "pytest>=7.0",
33
33
  "pytest-cov>=3.0",
34
34
  "black>=22.0",
35
- "mypy>=0.900",
36
- "ruff>=0.0.280",
37
- "pre-commit>=3.0",
38
35
  "build>=1.3",
39
36
  "scikit-learn>=1.0",
37
+ "yaml-changelog>=0.3.0",
38
+ "twine>=4.0.0",
40
39
  ]
41
40
  docs = [
42
41
  "mystmd>=1.3",
@@ -62,22 +61,6 @@ Issues = "https://github.com/PolicyEngine/L0/issues"
62
61
  line-length = 79
63
62
  target-version = ['py313']
64
63
 
65
- [tool.ruff]
66
- line-length = 79
67
- target-version = "py313"
68
- extend-exclude = ["docs/examples/*.ipynb"]
69
-
70
- [tool.ruff.lint]
71
- select = ["E", "F", "W", "I", "N", "UP"]
72
- ignore = ["E501", "N812", "UP038"] # E501: Line length handled by black, N812: Allow F as alias for functional, UP038: Use tuple for isinstance
73
-
74
- [tool.mypy]
75
- python_version = "3.13"
76
- warn_return_any = true
77
- warn_unused_configs = true
78
- disallow_untyped_defs = true
79
- ignore_missing_imports = true
80
-
81
64
  [tool.pytest.ini_options]
82
65
  testpaths = ["tests"]
83
66
  python_files = ["test_*.py"]
@@ -0,0 +1,423 @@
1
+ """
2
+ Tests for SparseCalibrationWeights with positive weight constraints.
3
+ """
4
+
5
+ import numpy as np
6
+ import torch
7
+ import pytest
8
+ from scipy import sparse as sp
9
+ from l0.calibration import SparseCalibrationWeights
10
+
11
+
12
+ class TestSparseCalibrationWeights:
13
+ """Test suite for calibration weights with L0 sparsity."""
14
+
15
+ def test_positive_weights(self):
16
+ """Verify all weights remain non-negative."""
17
+ N = 100
18
+ Q = 20
19
+
20
+ # Create test data
21
+ M = sp.random(Q, N, density=0.3, format="csr")
22
+ y = np.random.randn(Q) + 10
23
+
24
+ model = SparseCalibrationWeights(n_features=N)
25
+ model.fit(M, y, epochs=100, verbose=False)
26
+
27
+ # Check positivity
28
+ with torch.no_grad():
29
+ weights = model.get_weights(deterministic=True)
30
+ assert torch.all(weights >= 0), "Weights must be non-negative"
31
+
32
+ def test_sparse_ground_truth_relative_loss(self):
33
+ """Test recovery of sparse ground truth using relative loss."""
34
+ Q = 200 # targets
35
+ N = 2000 # features
36
+ N_active = 1000 # 50% sparsity
37
+
38
+ np.random.seed(42)
39
+ torch.manual_seed(42)
40
+
41
+ # Generate data with sparse ground truth
42
+ M_dense = np.random.lognormal(mean=1.5, sigma=0.25, size=(Q, N))
43
+ M = sp.csr_matrix(M_dense)
44
+
45
+ w_true = np.zeros(N)
46
+ active_indices = np.random.choice(N, size=N_active, replace=False)
47
+ w_true[active_indices] = np.random.lognormal(
48
+ mean=2.0, sigma=1.0, size=N_active
49
+ )
50
+
51
+ y = M @ w_true
52
+
53
+ # Fit with relative loss
54
+ model = SparseCalibrationWeights(
55
+ n_features=N,
56
+ beta=0.66,
57
+ gamma=-0.1,
58
+ zeta=1.1,
59
+ init_keep_prob=0.3,
60
+ init_weight_scale=0.5,
61
+ )
62
+
63
+ model.fit(
64
+ M=M,
65
+ y=y,
66
+ lambda_l0=0.0005, # Tuned for ~50% sparsity with relative loss
67
+ lambda_l2=1e-6,
68
+ lr=0.2,
69
+ epochs=2000,
70
+ loss_type="relative",
71
+ verbose=False,
72
+ )
73
+
74
+ # Check sparsity is reasonable (between 30% and 70%)
75
+ sparsity = model.get_sparsity()
76
+ assert (
77
+ 0.3 <= sparsity <= 0.7
78
+ ), f"Sparsity {sparsity:.2%} not in expected range"
79
+
80
+ # Check relative loss is low
81
+ with torch.no_grad():
82
+ y_pred = model.predict(M).cpu().numpy()
83
+ rel_loss = np.mean(((y - y_pred) / (y + 1)) ** 2)
84
+ assert rel_loss < 0.1, f"Relative loss {rel_loss:.4f} too high"
85
+
86
+ def test_relative_vs_mse_loss(self):
87
+ """Compare relative loss vs MSE for large-scale data."""
88
+ Q = 100
89
+ N = 500
90
+
91
+ np.random.seed(123)
92
+ torch.manual_seed(123)
93
+
94
+ # Large-scale data
95
+ M = sp.random(Q, N, density=0.5, format="csr")
96
+ M.data = np.abs(M.data) * 1000 # Large values
97
+ y = np.random.uniform(1000, 100000, size=Q)
98
+
99
+ # Train with MSE
100
+ model_mse = SparseCalibrationWeights(n_features=N)
101
+ model_mse.fit(
102
+ M,
103
+ y,
104
+ lambda_l0=1e-10, # Very small for MSE
105
+ lr=0.1,
106
+ epochs=500,
107
+ loss_type="mse",
108
+ verbose=False,
109
+ )
110
+
111
+ # Train with relative loss
112
+ model_rel = SparseCalibrationWeights(n_features=N)
113
+ model_rel.fit(
114
+ M,
115
+ y,
116
+ lambda_l0=0.001, # Can use larger penalty
117
+ lr=0.1,
118
+ epochs=500,
119
+ loss_type="relative",
120
+ verbose=False,
121
+ )
122
+
123
+ # Relative loss should achieve better relative accuracy
124
+ with torch.no_grad():
125
+ y_pred_mse = model_mse.predict(M).cpu().numpy()
126
+ y_pred_rel = model_rel.predict(M).cpu().numpy()
127
+
128
+ rel_err_mse = np.mean(np.abs((y - y_pred_mse) / (y + 1)))
129
+ rel_err_rel = np.mean(np.abs((y - y_pred_rel) / (y + 1)))
130
+
131
+ # Relative loss should do better on relative error
132
+ assert (
133
+ rel_err_rel <= rel_err_mse * 1.5
134
+ ), f"Relative loss should handle scale better: {rel_err_rel:.4f} vs {rel_err_mse:.4f}"
135
+
136
+ def test_sparsity_control(self):
137
+ """Test that L0 penalty controls sparsity level."""
138
+ Q = 50
139
+ N = 200
140
+
141
+ np.random.seed(123)
142
+ torch.manual_seed(123)
143
+
144
+ M = sp.random(Q, N, density=0.3, format="csr")
145
+ y = np.random.randn(Q) + 10
146
+
147
+ sparsities = []
148
+
149
+ # Test different L0 penalties
150
+ for lambda_l0 in [0.0001, 0.001, 0.01]:
151
+ model = SparseCalibrationWeights(n_features=N, init_keep_prob=0.5)
152
+ model.fit(
153
+ M,
154
+ y,
155
+ lambda_l0=lambda_l0,
156
+ lr=0.1,
157
+ epochs=2000,
158
+ loss_type="relative",
159
+ verbose=False,
160
+ )
161
+ sparsities.append(model.get_sparsity())
162
+
163
+ # Higher penalty should give more sparsity
164
+ assert (
165
+ sparsities[0] < sparsities[1]
166
+ ), "Higher L0 penalty should increase sparsity"
167
+ assert (
168
+ sparsities[1] < sparsities[2]
169
+ ), "Higher L0 penalty should increase sparsity"
170
+
171
+ def test_get_active_weights(self):
172
+ """Test active weight extraction."""
173
+ N = 100
174
+ model = SparseCalibrationWeights(n_features=N)
175
+
176
+ # Simple test data
177
+ M = sp.eye(N, format="csr")
178
+ y = np.ones(N)
179
+
180
+ model.fit(M, y, lambda_l0=0.01, epochs=100, verbose=False)
181
+
182
+ active_info = model.get_active_weights()
183
+
184
+ assert "indices" in active_info
185
+ assert "values" in active_info
186
+ assert "count" in active_info
187
+ assert active_info["count"] == len(active_info["indices"])
188
+ assert len(active_info["values"]) == active_info["count"]
189
+
190
+ # All active values should be positive
191
+ if active_info["count"] > 0:
192
+ assert torch.all(active_info["values"] > 0)
193
+
194
+ def test_deterministic_inference(self):
195
+ """Test that inference is deterministic."""
196
+ N = 50
197
+ Q = 10
198
+
199
+ np.random.seed(123)
200
+ torch.manual_seed(123)
201
+
202
+ M = sp.random(Q, N, density=0.5, format="csr")
203
+ y = np.random.randn(Q)
204
+
205
+ model = SparseCalibrationWeights(n_features=N)
206
+ model.fit(M, y, epochs=100, verbose=False)
207
+
208
+ # Multiple predictions should be identical
209
+ with torch.no_grad():
210
+ pred1 = model.predict(M).cpu().numpy()
211
+ pred2 = model.predict(M).cpu().numpy()
212
+
213
+ np.testing.assert_array_equal(
214
+ pred1, pred2, "Predictions should be deterministic"
215
+ )
216
+
217
+ def test_l2_regularization(self):
218
+ """Test that L2 penalty prevents weight explosion."""
219
+ N = 100
220
+ Q = 20
221
+
222
+ np.random.seed(123)
223
+ torch.manual_seed(123)
224
+
225
+ M = sp.random(Q, N, density=0.3, format="csr")
226
+ y = np.random.randn(Q) * 100 # Large scale
227
+
228
+ # Train without L2
229
+ model_no_l2 = SparseCalibrationWeights(n_features=N)
230
+ model_no_l2.fit(
231
+ M, y, lambda_l0=0.0001, lambda_l2=0.0, epochs=200, verbose=False
232
+ )
233
+
234
+ # Train with L2
235
+ model_with_l2 = SparseCalibrationWeights(n_features=N)
236
+ model_with_l2.fit(
237
+ M, y, lambda_l0=0.0001, lambda_l2=0.01, epochs=200, verbose=False
238
+ )
239
+
240
+ with torch.no_grad():
241
+ weights_no_l2 = model_no_l2.get_weights(deterministic=True)
242
+ weights_with_l2 = model_with_l2.get_weights(deterministic=True)
243
+
244
+ # L2 should reduce weight magnitudes
245
+ assert (
246
+ weights_with_l2.max() <= weights_no_l2.max() * 2.0
247
+ ), "L2 should prevent extreme weights"
248
+
249
+ def test_group_wise_averaging(self):
250
+ """Test that group-wise averaging balances loss contributions."""
251
+ N = 100 # features (households)
252
+
253
+ # Create targets with different cardinalities:
254
+ # - 3 singleton targets (like national targets)
255
+ # - 18 targets in one group (like age bins for one state)
256
+ # - 18 targets in another group (like age bins for another state)
257
+ Q = 3 + 18 + 18 # 39 total targets
258
+
259
+ np.random.seed(42)
260
+ torch.manual_seed(42)
261
+
262
+ # Create matrix with varying scales
263
+ M = sp.random(Q, N, density=0.3, format="csr")
264
+
265
+ # Create target values with different scales
266
+ # Singletons: large values (billions scale)
267
+ y_singletons = np.array([1e9, 5e8, 2e9])
268
+ # Groups: smaller values (thousands scale)
269
+ y_group1 = np.random.uniform(1e3, 1e6, size=18)
270
+ y_group2 = np.random.uniform(1e3, 1e6, size=18)
271
+ y = np.concatenate([y_singletons, y_group1, y_group2])
272
+
273
+ # Create target groups
274
+ # Groups 0, 1, 2: singletons (each national target)
275
+ # Group 3: all 18 targets from first age group
276
+ # Group 4: all 18 targets from second age group
277
+ target_groups = np.array(
278
+ [0, 1, 2] # 3 singletons
279
+ + [3] * 18 # Group 3
280
+ + [4] * 18 # Group 4
281
+ )
282
+
283
+ # Train WITHOUT grouping (baseline)
284
+ model_no_groups = SparseCalibrationWeights(n_features=N)
285
+ model_no_groups.fit(
286
+ M,
287
+ y,
288
+ lambda_l0=0.0001,
289
+ lr=0.1,
290
+ epochs=500,
291
+ loss_type="relative",
292
+ verbose=False,
293
+ target_groups=None, # No grouping
294
+ )
295
+
296
+ # Train WITH grouping
297
+ model_with_groups = SparseCalibrationWeights(n_features=N)
298
+ model_with_groups.fit(
299
+ M,
300
+ y,
301
+ lambda_l0=0.0001,
302
+ lr=0.1,
303
+ epochs=500,
304
+ loss_type="relative",
305
+ verbose=False,
306
+ target_groups=target_groups,
307
+ )
308
+
309
+ # Compute errors by group
310
+ with torch.no_grad():
311
+ y_pred_no_groups = model_no_groups.predict(M).cpu().numpy()
312
+ y_pred_with_groups = model_with_groups.predict(M).cpu().numpy()
313
+
314
+ # Relative errors
315
+ rel_err_no_groups = np.abs((y - y_pred_no_groups) / (y + 1))
316
+ rel_err_with_groups = np.abs((y - y_pred_with_groups) / (y + 1))
317
+
318
+ # Average errors by group
319
+ singleton_err_no_groups = rel_err_no_groups[:3].mean()
320
+ group3_err_no_groups = rel_err_no_groups[3:21].mean()
321
+ group4_err_no_groups = rel_err_no_groups[21:].mean()
322
+
323
+ singleton_err_with_groups = rel_err_with_groups[:3].mean()
324
+ group3_err_with_groups = rel_err_with_groups[3:21].mean()
325
+ group4_err_with_groups = rel_err_with_groups[21:].mean()
326
+
327
+ # With grouping, singleton errors should be much better
328
+ # (they're not dominated by the 36 histogram targets)
329
+ assert singleton_err_with_groups < singleton_err_no_groups * 1.5, (
330
+ f"Grouping should improve singleton accuracy: "
331
+ f"{singleton_err_with_groups:.4f} vs {singleton_err_no_groups:.4f}"
332
+ )
333
+
334
+ # All groups should have relatively balanced errors with grouping
335
+ all_group_errors = [
336
+ singleton_err_with_groups,
337
+ group3_err_with_groups,
338
+ group4_err_with_groups,
339
+ ]
340
+ max_err = max(all_group_errors)
341
+ min_err = min(all_group_errors)
342
+
343
+ # Errors should be within an order of magnitude of each other
344
+ assert max_err < min_err * 10, (
345
+ f"Group errors should be balanced: "
346
+ f"min={min_err:.4f}, max={max_err:.4f}"
347
+ )
348
+
349
+ def test_group_wise_averaging_edge_cases(self):
350
+ """Test edge cases for group-wise averaging."""
351
+ N = 50
352
+ Q = 10
353
+
354
+ np.random.seed(42)
355
+ torch.manual_seed(42)
356
+
357
+ M = sp.random(Q, N, density=0.3, format="csr")
358
+ y = np.random.uniform(100, 1000, size=Q)
359
+
360
+ model = SparseCalibrationWeights(n_features=N)
361
+
362
+ # Test 1: All targets in one group (should behave like no grouping)
363
+ target_groups_single = np.zeros(Q, dtype=int)
364
+ model.fit(
365
+ M,
366
+ y,
367
+ lambda_l0=0.00001, # Lower penalty for better convergence
368
+ epochs=2000, # Plenty of epochs
369
+ lr=0.2, # Higher learning rate
370
+ loss_type="relative",
371
+ verbose=False,
372
+ target_groups=target_groups_single,
373
+ )
374
+
375
+ with torch.no_grad():
376
+ y_pred = model.predict(M).cpu().numpy()
377
+ rel_err = np.mean(np.abs((y - y_pred) / (y + 1)))
378
+ assert (
379
+ rel_err < 0.5
380
+ ), f"Single group should still converge, got {rel_err:.4f}"
381
+
382
+ # Test 2: Each target in its own group (like all singletons)
383
+ target_groups_all_singleton = np.arange(Q)
384
+ model_new = SparseCalibrationWeights(n_features=N)
385
+ model_new.fit(
386
+ M,
387
+ y,
388
+ lambda_l0=0.00001,
389
+ epochs=2000,
390
+ lr=0.2,
391
+ loss_type="relative",
392
+ verbose=False,
393
+ target_groups=target_groups_all_singleton,
394
+ )
395
+
396
+ with torch.no_grad():
397
+ y_pred = model_new.predict(M).cpu().numpy()
398
+ rel_err = np.mean(np.abs((y - y_pred) / (y + 1)))
399
+ assert (
400
+ rel_err < 0.5
401
+ ), f"All singleton groups should converge, got {rel_err:.4f}"
402
+
403
+ # Test 3: Unbalanced groups (1 huge group, several small)
404
+ target_groups_unbalanced = np.array([0] * 7 + [1, 2, 3])
405
+ model_unbalanced = SparseCalibrationWeights(n_features=N)
406
+ model_unbalanced.fit(
407
+ M,
408
+ y,
409
+ lambda_l0=0.00001,
410
+ epochs=2000,
411
+ lr=0.2,
412
+ loss_type="relative",
413
+ verbose=False,
414
+ target_groups=target_groups_unbalanced,
415
+ )
416
+
417
+ with torch.no_grad():
418
+ y_pred = model_unbalanced.predict(M).cpu().numpy()
419
+ # Check that small groups aren't ignored
420
+ small_group_errors = np.abs((y[7:] - y_pred[7:]) / (y[7:] + 1))
421
+ assert (
422
+ np.mean(small_group_errors) < 0.5
423
+ ), "Small groups should not be ignored"
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes