l0-python 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

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.
l0/__init__.py CHANGED
@@ -26,6 +26,7 @@ from .penalties import (
26
26
  get_sparsity_stats,
27
27
  update_temperatures,
28
28
  )
29
+ from .sparse import SparseL0Linear
29
30
 
30
31
  __all__ = [
31
32
  # Distributions
@@ -36,6 +37,8 @@ __all__ = [
36
37
  "L0DepthwiseConv2d",
37
38
  "SparseMLP",
38
39
  "prune_model",
40
+ # Sparse
41
+ "SparseL0Linear",
39
42
  # Gates
40
43
  "L0Gate",
41
44
  "SampleGate",
l0/calibration.py ADDED
@@ -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)
l0/sparse.py ADDED
@@ -0,0 +1,321 @@
1
+ """
2
+ L0 regularized linear models for sparse data.
3
+
4
+ Implements L0-regularized linear regression that efficiently handles
5
+ sparse input matrices without converting them to dense format.
6
+ """
7
+
8
+ # ruff: noqa: N803 N806
9
+
10
+
11
+ import numpy as np
12
+ import torch
13
+ import torch.nn as nn
14
+ from scipy import sparse as sp
15
+
16
+
17
+ class SparseL0Linear(nn.Module):
18
+ """
19
+ L0-regularized linear model for sparse input data.
20
+
21
+ Efficiently handles sparse matrices (scipy.sparse) without
22
+ converting to dense format during computation.
23
+
24
+ Parameters
25
+ ----------
26
+ n_features : int
27
+ Number of input features
28
+ fit_intercept : bool
29
+ Whether to fit an intercept term
30
+ beta : float
31
+ Temperature parameter for Hard Concrete distribution
32
+ gamma : float
33
+ Lower bound of stretched concrete distribution
34
+ zeta : float
35
+ Upper bound of stretched concrete distribution
36
+ init_keep_prob : float
37
+ Initial probability of keeping each feature
38
+ device : str or torch.device
39
+ Device to run computations on ('cpu' or 'cuda')
40
+ """
41
+
42
+ def __init__(
43
+ self,
44
+ n_features: int,
45
+ fit_intercept: bool = True,
46
+ beta: float = 2 / 3,
47
+ gamma: float = -0.1,
48
+ zeta: float = 1.1,
49
+ init_keep_prob: float = 0.5,
50
+ device: str | torch.device = "cpu",
51
+ ):
52
+ super().__init__()
53
+ self.n_features = n_features
54
+ self.fit_intercept = fit_intercept
55
+ self.beta = beta
56
+ self.gamma = gamma
57
+ self.zeta = zeta
58
+ self.device = torch.device(device)
59
+
60
+ # Model parameters
61
+ self.weight = nn.Parameter(torch.zeros(n_features, device=self.device))
62
+ if fit_intercept:
63
+ self.bias = nn.Parameter(torch.zeros(1, device=self.device))
64
+ else:
65
+ self.register_parameter("bias", None)
66
+
67
+ # L0 gate parameters
68
+ mu = torch.log(torch.tensor(init_keep_prob / (1 - init_keep_prob)))
69
+ self.log_alpha = nn.Parameter(
70
+ torch.normal(
71
+ mu.item(), 0.01, size=(n_features,), device=self.device
72
+ )
73
+ )
74
+
75
+ # Cache for sparse tensor conversion
76
+ self._cached_X_torch: torch.sparse.Tensor | None = None
77
+ self._cached_X_shape: tuple[int, int] | None = None
78
+
79
+ def _convert_sparse_to_torch(
80
+ self, X_sparse: sp.spmatrix
81
+ ) -> torch.sparse.Tensor:
82
+ """
83
+ Convert scipy sparse matrix to torch sparse tensor.
84
+
85
+ Caches the result if the shape matches to avoid redundant conversions.
86
+ """
87
+ if (
88
+ self._cached_X_torch is not None
89
+ and self._cached_X_shape == X_sparse.shape
90
+ and X_sparse.nnz == self._cached_X_torch._nnz()
91
+ ):
92
+ return self._cached_X_torch
93
+
94
+ X_coo = X_sparse.tocoo()
95
+ indices = torch.LongTensor(np.vstack([X_coo.row, X_coo.col])).to(
96
+ self.device
97
+ )
98
+ values = torch.FloatTensor(X_coo.data).to(self.device)
99
+ X_torch = torch.sparse_coo_tensor(
100
+ indices,
101
+ values,
102
+ X_sparse.shape,
103
+ dtype=torch.float32,
104
+ device=self.device,
105
+ )
106
+
107
+ # Cache for future use
108
+ self._cached_X_torch = X_torch
109
+ self._cached_X_shape = X_sparse.shape
110
+
111
+ return X_torch
112
+
113
+ def _sample_gates(self) -> torch.Tensor:
114
+ """Sample gates using Hard Concrete distribution."""
115
+ eps = 1e-6
116
+ u = torch.rand_like(self.log_alpha).clamp(eps, 1 - eps)
117
+ X = (torch.log(u) - torch.log(1 - u) + self.log_alpha) / self.beta
118
+ s = torch.sigmoid(X)
119
+ s_bar = s * (self.zeta - self.gamma) + self.gamma
120
+ return s_bar.clamp(0, 1)
121
+
122
+ def get_deterministic_gates(self) -> torch.Tensor:
123
+ """Get deterministic gate values (for inference)."""
124
+ X = self.log_alpha / self.beta
125
+ s = torch.sigmoid(X)
126
+ s_bar = s * (self.zeta - self.gamma) + self.gamma
127
+ return s_bar.clamp(0, 1)
128
+
129
+ def forward(
130
+ self,
131
+ X: torch.sparse.Tensor | sp.spmatrix,
132
+ deterministic: bool = False,
133
+ ) -> torch.Tensor:
134
+ """
135
+ Forward pass maintaining sparse computation.
136
+
137
+ Parameters
138
+ ----------
139
+ X : torch.sparse.Tensor or scipy.sparse matrix
140
+ Sparse input tensor of shape (n_samples, n_features)
141
+ deterministic : bool
142
+ Whether to use deterministic gates (for inference)
143
+
144
+ Returns
145
+ -------
146
+ torch.Tensor
147
+ Predictions of shape (n_samples,)
148
+ """
149
+ # Convert scipy sparse to torch if needed
150
+ if isinstance(X, sp.spmatrix):
151
+ X = self._convert_sparse_to_torch(X)
152
+
153
+ # Sample or get deterministic gates
154
+ if deterministic:
155
+ gates = self.get_deterministic_gates()
156
+ else:
157
+ gates = self._sample_gates()
158
+
159
+ # Apply gates to weights
160
+ gated_weight = self.weight * gates
161
+
162
+ # Sparse matrix multiplication
163
+ # X is (n_samples, n_features), gated_weight is (n_features,)
164
+ y = torch.sparse.mm(X, gated_weight.unsqueeze(1)).squeeze(1)
165
+
166
+ # Add bias if needed
167
+ if self.fit_intercept and self.bias is not None:
168
+ y = y + self.bias
169
+
170
+ return y # type: ignore[no-any-return]
171
+
172
+ def get_l0_penalty(self) -> torch.Tensor:
173
+ """
174
+ Compute L0 complexity penalty.
175
+
176
+ Returns expected number of active parameters.
177
+ """
178
+ c = -self.beta * torch.log(
179
+ torch.tensor(-self.gamma / self.zeta, device=self.device)
180
+ )
181
+ pi = torch.sigmoid(self.log_alpha + c)
182
+ return pi.sum()
183
+
184
+ def get_sparsity(self) -> float:
185
+ """
186
+ Get current sparsity level.
187
+
188
+ Returns
189
+ -------
190
+ float
191
+ Fraction of weights that are effectively zero
192
+ """
193
+ with torch.no_grad():
194
+ gates = self.get_deterministic_gates()
195
+ return (gates < 0.01).float().mean().item()
196
+
197
+ def get_selected_features(self, threshold: float = 0.01) -> torch.Tensor:
198
+ """
199
+ Get indices of selected (non-zero) features.
200
+
201
+ Parameters
202
+ ----------
203
+ threshold : float
204
+ Gate values below this are considered zero
205
+
206
+ Returns
207
+ -------
208
+ torch.Tensor
209
+ Indices of selected features
210
+ """
211
+ with torch.no_grad():
212
+ gates = self.get_deterministic_gates()
213
+ return torch.where(gates > threshold)[0]
214
+
215
+ def fit(
216
+ self,
217
+ X_sparse: sp.spmatrix,
218
+ y: torch.Tensor | np.ndarray,
219
+ lambda_reg: float = 0.01,
220
+ lr: float = 0.01,
221
+ epochs: int = 1000,
222
+ verbose: bool = False,
223
+ verbose_freq: int = 100,
224
+ ) -> "SparseL0Linear":
225
+ """
226
+ Fit the model using gradient descent.
227
+
228
+ Parameters
229
+ ----------
230
+ X_sparse : scipy.sparse matrix
231
+ Sparse input features
232
+ y : torch.Tensor or numpy.ndarray
233
+ Target values
234
+ lambda_reg : float
235
+ L0 regularization strength
236
+ lr : float
237
+ Learning rate
238
+ epochs : int
239
+ Number of training epochs
240
+ verbose : bool
241
+ Whether to print progress
242
+ verbose_freq : int
243
+ How often to print progress
244
+
245
+ Returns
246
+ -------
247
+ self
248
+ Fitted model
249
+ """
250
+ import numpy as np
251
+
252
+ # Convert y to tensor if needed
253
+ if isinstance(y, np.ndarray):
254
+ y = torch.tensor(y, dtype=torch.float32, device=self.device)
255
+ elif not isinstance(y, torch.Tensor):
256
+ y = torch.tensor(y, dtype=torch.float32, device=self.device)
257
+ else:
258
+ y = y.to(self.device)
259
+
260
+ # Convert sparse matrix to torch sparse tensor (will be cached)
261
+ X_torch = self._convert_sparse_to_torch(X_sparse)
262
+
263
+ # Initialize weights with small random values
264
+ nn.init.normal_(self.weight, 0, 0.1)
265
+ if self.fit_intercept and self.bias is not None:
266
+ self.bias.data.fill_(y.mean())
267
+
268
+ # Setup optimizer
269
+ params = [self.weight, self.log_alpha]
270
+ if self.fit_intercept and self.bias is not None:
271
+ params.append(self.bias)
272
+ optimizer = torch.optim.Adam(params, lr=lr)
273
+
274
+ # Training loop
275
+ for epoch in range(epochs):
276
+ # Forward pass
277
+ y_pred = self.forward(X_torch, deterministic=False)
278
+
279
+ # Compute loss
280
+ mse_loss = (y - y_pred).pow(2).mean()
281
+ l0_loss = self.get_l0_penalty()
282
+ loss = mse_loss + lambda_reg * l0_loss
283
+
284
+ # Backward pass
285
+ optimizer.zero_grad()
286
+ loss.backward()
287
+ optimizer.step()
288
+
289
+ # Print progress
290
+ if verbose and (epoch + 1) % verbose_freq == 0:
291
+ with torch.no_grad():
292
+ n_active = (self.get_deterministic_gates() > 0.01).sum()
293
+ print(
294
+ f"Epoch {epoch+1:4d}: "
295
+ f"loss={loss.item():.4f}, "
296
+ f"mse={mse_loss.item():.4f}, "
297
+ f"l0={l0_loss.item():.2f}, "
298
+ f"active={n_active.item()}"
299
+ )
300
+
301
+ return self
302
+
303
+ def predict(
304
+ self, X_sparse: sp.spmatrix | torch.sparse.Tensor
305
+ ) -> torch.Tensor:
306
+ """
307
+ Make predictions on new data.
308
+
309
+ Parameters
310
+ ----------
311
+ X_sparse : scipy.sparse matrix or torch.sparse.Tensor
312
+ Sparse input features
313
+
314
+ Returns
315
+ -------
316
+ torch.Tensor
317
+ Predictions
318
+ """
319
+ # Make deterministic predictions
320
+ with torch.no_grad():
321
+ return self.forward(X_sparse, deterministic=True)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: l0-python
3
- Version: 0.1.0
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
@@ -41,14 +41,15 @@ Description-Content-Type: text/markdown
41
41
  License-File: LICENSE
42
42
  Requires-Dist: torch>=2.0.0
43
43
  Requires-Dist: numpy>=1.20.0
44
+ Requires-Dist: scipy>=1.7.0
44
45
  Provides-Extra: dev
45
46
  Requires-Dist: pytest>=7.0; extra == "dev"
46
47
  Requires-Dist: pytest-cov>=3.0; extra == "dev"
47
48
  Requires-Dist: black>=22.0; extra == "dev"
48
- Requires-Dist: mypy>=0.900; extra == "dev"
49
- Requires-Dist: ruff>=0.0.280; extra == "dev"
50
- Requires-Dist: pre-commit>=3.0; extra == "dev"
51
49
  Requires-Dist: build>=1.3; extra == "dev"
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"
52
53
  Provides-Extra: docs
53
54
  Requires-Dist: mystmd>=1.3; extra == "docs"
54
55
  Requires-Dist: jupyter>=1.0; extra == "docs"
@@ -0,0 +1,12 @@
1
+ l0/__init__.py,sha256=ezGVLGbdU_jAElAJSdLKjSWCnZRILAZi8qMBPjcHwQc,1192
2
+ l0/calibration.py,sha256=H-5bo0SOCQXNBFuAQCzfBxw9-j4Ta0DcmJ66Hv-3Qtc,14860
3
+ l0/distributions.py,sha256=GvVoDasdxoDlPPdzt33lFzkuJgQsfkEn9e-abxru3nM,6325
4
+ l0/gates.py,sha256=54OORMg9qSsdbIyzYIL9t9CCHWsKbumSKl_d4-YFuNU,11443
5
+ l0/layers.py,sha256=j_DPEFdfUc37ojtoK2ufGzVYzZ_DwusqTcJXCVxZS68,13785
6
+ l0/penalties.py,sha256=okqXgX2eGUmg5prcDclicTQALPPPF9FA5eQP09V4GTE,9569
7
+ l0/sparse.py,sha256=UVpkrrJBnQ9iSNDoWK5F_sr43tQareWXzYHyMq_cKss,9726
8
+ l0_python-0.3.0.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
9
+ l0_python-0.3.0.dist-info/METADATA,sha256=sBDt2EahzEyNWSBsufl0geYWY5mwq9JMkSTVGppjPkg,6833
10
+ l0_python-0.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
11
+ l0_python-0.3.0.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
12
+ l0_python-0.3.0.dist-info/RECORD,,
@@ -1,10 +0,0 @@
1
- l0/__init__.py,sha256=meL6jJ1CxO6pO373nY-nwGXnpzWZGC1Ts6q2TLyAeLM,1122
2
- l0/distributions.py,sha256=GvVoDasdxoDlPPdzt33lFzkuJgQsfkEn9e-abxru3nM,6325
3
- l0/gates.py,sha256=54OORMg9qSsdbIyzYIL9t9CCHWsKbumSKl_d4-YFuNU,11443
4
- l0/layers.py,sha256=j_DPEFdfUc37ojtoK2ufGzVYzZ_DwusqTcJXCVxZS68,13785
5
- l0/penalties.py,sha256=okqXgX2eGUmg5prcDclicTQALPPPF9FA5eQP09V4GTE,9569
6
- l0_python-0.1.0.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
7
- l0_python-0.1.0.dist-info/METADATA,sha256=hrAugZ-F7oKmleyJ8WjVTykRixBmuTiQ2M4oae1Ceu4,6794
8
- l0_python-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
9
- l0_python-0.1.0.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
10
- l0_python-0.1.0.dist-info/RECORD,,