Skip to content

learnable_regularizers

Learnable filters for custom decompositions.

Classes:

Functions:

  • estimate_lambdas –

    Estimate per-filter thresholds lambda_i after sparsity-based training.

  • huber_moreau_prox –

    Proximal operator of sigma * (lamphi_delta)^ via Moreau decomposition:

  • huber_prox –

    Proximal operator of lam * phi_delta applied element-wise.

  • train_lambdas_denoising –

    Learn the per-filter thresholds/weights lambda_i by minimizing the MSE denoising loss:

LearnableRegularizerHuber 

LearnableRegularizerHuber(
    filterbank: ConvolutionalDecompositionBase,
    lambda_init: float | NDArray | Tensor = 0.05,
    delta_init: float | NDArray | Tensor = 0.1,
    delta_floor: float = eps,
    fix_dc: bool = True,
)

Bases: LearnableRegularizerL1

R(x) = \sum_i \sum_n lambda_i * phi_{delta_i}((q_i * x)[n])

with {q_i} forming a Parseval filterbank (T T^T = I).

Learnable parameters: log_lam : (1, m, 1, 1) - log(lambda_i), all channels log_delta : (1, m, 1, 1) - log(delta_i), all channels

The DC/constant filter (channel 0) is handled by fix_dc=True, which sets lambda_0 = 0 (no regularization of the mean coefficient).

Parameters:

Methods:

  • dual_prox –

    prox_{sigma * (scale * R)^T}(u) = clamp(u, -scale * lambda, scale * lambda)

  • evaluate –

    R(x) per batch element, shape (B,).

  • prox –

    prox_{scale*R}(v) in IMAGE domain.

  • trainable_params –

    Return only the genuinely trained parameters (excluding frozen fb).

Attributes:

  • deltas (Tensor) –

    delta_i > 0, shape (1, m, 1, 1).

  • lambdas (Tensor) –

    lambda_i > 0, shape (1, m, 1, 1). DC channel forced to 0 if fix_dc.

Source code in src/autoden/transforms/learnable_regularizers.py 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
def __init__(
    self,
    filterbank: ConvolutionalDecompositionBase,
    lambda_init: float | NDArray | pt.Tensor = 0.05,
    delta_init: float | NDArray | pt.Tensor = 0.1,
    delta_floor: float = eps,
    fix_dc: bool = True,
):
    super().__init__(filterbank=filterbank, lambda_init=lambda_init, fix_dc=fix_dc)
    m = filterbank.m
    ones_k = (1,) * filterbank.n_dims

    if isinstance(delta_init, float):
        delta_init = pt.full((1, m, *ones_k), delta_init)
    elif isinstance(delta_init, (np.ndarray, pt.Tensor)):
        if isinstance(delta_init, np.ndarray):
            delta_init = pt.tensor(delta_init)
        if delta_init.numel() != m:
            raise ValueError(f"The number of `delta_init` should be: {m}, but {delta_init.numel()} found instead.")
        kernels = filterbank.get_kernels()
        delta_init = delta_init.to(kernels.device, dtype=kernels.dtype).view((1, m, *ones_k))
    else:
        raise ValueError("Parameter `init_lambda` should be one of: float | NDArray | pt.Tensor")
    self.log_delta = nn.Parameter(delta_init)
    self.delta_floor = delta_floor

deltas property 

deltas: Tensor

delta_i > 0, shape (1, m, 1, 1).

lambdas property 

lambdas: Tensor

lambda_i > 0, shape (1, m, 1, 1). DC channel forced to 0 if fix_dc.

dual_prox 

dual_prox(u: Tensor, scale: float = 1.0) -> Tensor

prox_{sigma * (scale * R)^T}(u) = clamp(u, -scale * lambda, scale * lambda)

This is the Moreau-dual proximal, used in the PDHG dual update: u <- dual_prox(u + sigma * Wx, scale=1) with lambda already encoded or equivalently clamp(u + sigma * Wx, -lambda, lambda) for the problem: lambda * ||Wx||_1.

The clamp bound passed here is scale * lambda_i per channel i.

Source code in src/autoden/transforms/learnable_regularizers.py 
235
236
237
238
239
240
241
242
243
244
245
def dual_prox(self, u: pt.Tensor, scale: float = 1.0) -> pt.Tensor:
    """
    prox_{sigma * (scale * R)^T}(u) = clamp(u, -scale * lambda, scale * lambda)

    This is the Moreau-dual proximal, used in the PDHG dual update:
        u <- dual_prox(u + sigma * Wx,  scale=1)    with lambda already encoded
    or equivalently clamp(u + sigma * Wx, -lambda, lambda) for the problem: lambda * ||Wx||_1.

    The clamp bound passed here is scale * lambda_i per channel i.
    """
    return huber_moreau_prox(u, scale, self.lambdas, self.deltas)

evaluate 

evaluate(x: Tensor) -> Tensor

R(x) per batch element, shape (B,).

Source code in src/autoden/transforms/learnable_regularizers.py 
247
248
249
250
251
def evaluate(self, x: pt.Tensor) -> pt.Tensor:
    """R(x) per batch element, shape (B,)."""
    axes = [*range(1, self.fb.n_dims + 2)]
    v = self.fb.analyze(x)
    return huber_loss(v, self.lambdas, self.deltas).sum(dim=tuple(axes))

prox 

prox(v: Tensor, scale: float = 1.0) -> Tensor

prox_{scale*R}(v) in IMAGE domain.

For stride-k Parseval (T T^T = I): prox_{scaleR}(v) = T · prox_{scalePhi}(T^T v)

i.e. analyse into coefficient domain, apply per-channel Huber prox, synthesise back. Exact because T is an isometry.

Source code in src/autoden/transforms/learnable_regularizers.py 
220
221
222
223
224
225
226
227
228
229
230
231
232
def prox(self, v: pt.Tensor, scale: float = 1.0) -> pt.Tensor:
    """
    prox_{scale*R}(v) in IMAGE domain.

    For stride-k Parseval (T T^T = I):
        prox_{scale*R}(v) = T · prox_{scale*Phi}(T^T v)

    i.e. analyse into coefficient domain, apply per-channel Huber prox,
    synthesise back. Exact because T is an isometry.
    """
    Ttv = self.fb.analyze(v)  # (B, m, h, w)
    prxd = huber_prox(Ttv, self.lambdas * scale, self.deltas)
    return self.fb.synthesize(prxd)  # (B, C, H, W)

trainable_params 

trainable_params() -> list

Return only the genuinely trained parameters (excluding frozen fb).

Source code in src/autoden/transforms/learnable_regularizers.py 
215
216
217
def trainable_params(self) -> list:
    """Return only the genuinely trained parameters (excluding frozen fb)."""
    return [self.log_lambda, self.log_delta]

LearnableRegularizerL1 

LearnableRegularizerL1(
    filterbank: ConvolutionalDecompositionBase,
    lambda_init: float | NDArray | Tensor = 0.05,
    fix_dc: bool = True,
)

Bases: Module

R(x) = sum_i{ lambda_i ||q_i * x||_1 }

Proximal operator (ADMM / proximal-gradient style): prox_{scale * R}(v) = v - (1/m) * W^T * SoftThresh_{m * scale * lambda}(Wv)

Dual proximal operator (PDHG style, on the dual variable u): prox_{sigma * (scale * R)^T}(u) = clamp(u, -scale * lambda, scale * lambda)

Both are provided. The same lambda_i can be used for both algorithms, but the OPTIMAL lambda_i may differ because PDHG operates in the dual domain and the effective scale depends on the step sizes sigma, tau. Use train_lambdas_denoising_admm() or train_lambdas_denoising_pdhg() to calibrate for the desired algorithm.

Methods:

  • dual_prox –

    prox_{sigma * (scale * R)^T}(u) = clamp(u, -scale * lambda, scale * lambda)

  • evaluate –

    R(x) per batch element, shape (B,).

  • prox –

    prox_{scale * R}(v) = v - (1/m) * W^T * SoftThresh_{m * scale * lambda}(Wv)

  • trainable_params –

    Return only the genuinely trained parameters (excluding frozen fb).

Attributes:

  • lambdas (Tensor) –

    lambda_i > 0, shape (1, m, 1, 1). DC channel forced to 0 if fix_dc.

Source code in src/autoden/transforms/learnable_regularizers.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def __init__(
    self,
    filterbank: ConvolutionalDecompositionBase,
    lambda_init: float | NDArray | pt.Tensor = 0.05,
    fix_dc: bool = True,
):
    super().__init__()
    self.fb = filterbank
    self.fix_dc = fix_dc
    m = filterbank.m
    n_dims = filterbank.n_dims
    ones_k = (1,) * n_dims

    if isinstance(lambda_init, float):
        lambda_init = pt.full((1, m, *ones_k), lambda_init)
        if not fix_dc:
            lambda_init[0, 0] *= 0.01
    elif isinstance(lambda_init, (np.ndarray, pt.Tensor)):
        if isinstance(lambda_init, np.ndarray):
            lambda_init = pt.tensor(lambda_init)
        if lambda_init.numel() != m:
            raise ValueError(f"The number of `lambda_init` should be: {m}, but {lambda_init.numel()} found instead.")
        kernels = filterbank.get_kernels()
        lambda_init = lambda_init.to(kernels.device, dtype=kernels.dtype).view((1, m, *ones_k))
    else:
        raise ValueError("Parameter `init_lambda` should be one of: float | NDArray | pt.Tensor")
    self.log_lambda = nn.Parameter(lambda_init.log())

lambdas property 

lambdas: Tensor

lambda_i > 0, shape (1, m, 1, 1). DC channel forced to 0 if fix_dc.

dual_prox 

dual_prox(u: Tensor, scale: float = 1.0) -> Tensor

prox_{sigma * (scale * R)^T}(u) = clamp(u, -scale * lambda, scale * lambda)

This is the Moreau-dual proximal, used in the PDHG dual update: u <- dual_prox(u + sigma * Wx, scale=1) with lambda already encoded or equivalently clamp(u + sigma * Wx, -lambda, lambda) for the problem: lambda * ||Wx||_1.

The clamp bound passed here is scale * lambda_i per channel i.

Source code in src/autoden/transforms/learnable_regularizers.py 
103
104
105
106
107
108
109
110
111
112
113
def dual_prox(self, u: pt.Tensor, scale: float = 1.0) -> pt.Tensor:
    """
    prox_{sigma * (scale * R)^T}(u) = clamp(u, -scale * lambda, scale * lambda)

    This is the Moreau-dual proximal, used in the PDHG dual update:
        u <- dual_prox(u + sigma * Wx,  scale=1)    with lambda already encoded
    or equivalently clamp(u + sigma * Wx, -lambda, lambda) for the problem: lambda * ||Wx||_1.

    The clamp bound passed here is scale * lambda_i per channel i.
    """
    return pt.clamp(u, -self.lambdas * scale, self.lambdas * scale)

evaluate 

evaluate(x: Tensor) -> Tensor

R(x) per batch element, shape (B,).

Source code in src/autoden/transforms/learnable_regularizers.py 
115
116
117
118
def evaluate(self, x: pt.Tensor) -> pt.Tensor:
    """R(x) per batch element, shape (B,)."""
    axes = [*range(1, self.fb.n_dims + 2)]
    return (self.lambdas * self.fb.analyze(x).abs()).sum(dim=tuple(axes))

prox 

prox(v: Tensor, scale: float = 1.0) -> Tensor

prox_{scale * R}(v) = v - (1/m) * W^T * SoftThresh_{m * scale * lambda}(Wv)

Exact when FF^T = I_m AND spectral flatness holds (cond. A + B). Approximate (boundary only) when only cond. A holds.

Source code in src/autoden/transforms/learnable_regularizers.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def prox(self, v: pt.Tensor, scale: float = 1.0) -> pt.Tensor:
    """
    prox_{scale * R}(v) = v - (1/m) * W^T * SoftThresh_{m * scale * lambda}(Wv)

    Exact when FF^T = I_m AND spectral flatness holds (cond. A + B).
    Approximate (boundary only) when only cond. A holds.
    """
    Wv = self.fb.analyze(v)
    lam = self.lambdas * scale
    Wv_st = pt.sign(Wv) * F.relu(Wv.abs() - lam)
    shrinkage = Wv - Wv_st
    return v - self.fb.synthesize(shrinkage)

trainable_params 

trainable_params() -> list

Return only the genuinely trained parameters (excluding frozen fb).

Source code in src/autoden/transforms/learnable_regularizers.py 
84
85
86
def trainable_params(self) -> list:
    """Return only the genuinely trained parameters (excluding frozen fb)."""
    return [self.log_lambda]

estimate_lambdas 

estimate_lambdas(
    filterbank: ConvolutionalDecompositionBase,
    data_val: NDArray,
    sigma: float = 25.0 / 255.0,
    method: Literal["mad"] | Literal["sweep"] = "mad",
    lams: Sequence[float] | NDArray | None = None,
    filter_weights: NDArray | None = None,
    batch_size: int = 16,
    plot_result: bool = True,
    device: str = "cuda" if is_available() else "cpu",
) -> NDArray

Estimate per-filter thresholds lambda_i after sparsity-based training.

Parameters:

  • filterbank (ConvolutionalDecompositionBase) –

    The filterbank to estimate the thresholds for.

  • data_val (NDArray) –

    Clean images for calibration.

  • sigma (float, default: 25.0 / 255.0 ) –

    The standard deviation of the noise, by default 25.0 / 255.0.

  • method (Literal['mad'] | Literal['sweep'], default: 'mad' ) –

    The method to use for estimation, by default "mad".

  • lams (Sequence[float] | NDArray | None, default: None ) –

    The list of lambdas to test in the sweep method, by default None.

  • filter_weights (NDArray | None, default: None ) –

    Individual filter weights, by default None.

  • plot_result (bool, default: True ) –

    Whether to plot the results, by default True.

  • device (str, default: 'cuda' if is_available() else 'cpu' ) –

    The device to use for computation, by default "cuda" if available, else "cpu".

Returns:

  • NDArray –

    The estimated thresholds for each filter.

Notes

Two methods are available:

'mad' (Median Absolute Deviation / Donoho-Johnstone universal threshold): lambda_i = sigma_i * sqrt(2 * log(HW)) where sigma_i = MAD((W epsilon)_i) / 0.6745 estimates the noise std of filter i when applied to pure white noise epsilon ~ N(0, sigma^2 I). This is the classical wavelet denoising threshold.

'sweep': run a coarse grid search for the best lambda (single global lambda / or a global lambda multiplier) on the validation images. Fast, data-driven.

Source code in src/autoden/transforms/learnable_regularizers.py 
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
def estimate_lambdas(
    filterbank: ConvolutionalDecompositionBase,
    data_val: NDArray,  # clean images for calibration
    sigma: float = 25.0 / 255.0,
    method: Literal["mad"] | Literal["sweep"] = "mad",
    lams: Sequence[float] | NDArray | None = None,
    filter_weights: NDArray | None = None,
    batch_size: int = 16,
    plot_result: bool = True,
    device: str = "cuda" if pt.cuda.is_available() else "cpu",
) -> NDArray:
    """
    Estimate per-filter thresholds lambda_i after sparsity-based training.

    Parameters
    ----------
    filterbank : ConvolutionalDecompositionBase
        The filterbank to estimate the thresholds for.
    data_val : NDArray
        Clean images for calibration.
    sigma : float, optional
        The standard deviation of the noise, by default 25.0 / 255.0.
    method : Literal["mad"] | Literal["sweep"], optional
        The method to use for estimation, by default "mad".
    lams : Sequence[float] | NDArray | None, optional
        The list of lambdas to test in the sweep method, by default None.
    filter_weights : NDArray | None, optional
        Individual filter weights, by default None.
    plot_result : bool, optional
        Whether to plot the results, by default True.
    device : str, optional
        The device to use for computation, by default "cuda" if available, else "cpu".

    Returns
    -------
    NDArray
        The estimated thresholds for each filter.

    Notes
    -----
    Two methods are available:

    'mad' (Median Absolute Deviation / Donoho-Johnstone universal threshold):
        lambda_i = sigma_i * sqrt(2 * log(HW))
        where sigma_i = MAD((W epsilon)_i) / 0.6745 estimates the noise std of
        filter i when applied to pure white noise epsilon ~ N(0, sigma^2 I).
        This is the classical wavelet denoising threshold.

    'sweep': run a coarse grid search for the best lambda (single global lambda
        / or a global lambda multiplier) on the validation images. Fast, data-driven.
    """
    filterbank = filterbank.to(device)

    weights_shape: tuple[int, ...] = (1, filterbank.m, *(1,) * filterbank.n_dims)

    with pt.inference_mode():
        if method.lower() == "mad":
            # Estimate noise response of each filter
            noise = sigma * pt.rand_like(pt.from_numpy(data_val)).to(device)
            Wn = filterbank.analyze(noise)  # (N, m, H, W)
            # MAD per filter (robust std estimator)
            mad = Wn.abs().median(dim=0).values.median(dim=-1).values.median(dim=-1).values
            # Universal threshold: sigma_i * sqrt(2 log n)
            n = int(np.prod(data_val.shape[-filterbank.n_dims :]))
            lam = (mad / 0.6745) * math.sqrt(2 * math.log(n))
            return lam.view(*weights_shape).cpu().numpy().copy()

        elif method.lower() == "sweep":
            if lams is None:
                raise ValueError("Please provide a list of lambdas to test.")

            dset_val = DatasetNumpy(data_val, device=device)
            val_dsets_list = DatasetsList([dset_val, dset_val], augmentation=AugmentationGaussianNoise(sigma))
            val_dl = DataLoader(val_dsets_list, batch_size=batch_size, shuffle=False, num_workers=0)
            psnrs = np.zeros(len(lams))

            for ii, lam_val in enumerate(tqdm(lams, desc="Testing lambdas")):
                lam_t = pt.full(weights_shape, lam_val, device=device)
                if filter_weights is not None:
                    lam_t *= pt.tensor(filter_weights).to(device).reshape(weights_shape)

                psnr_sum = 0.0
                for noisy, x in val_dl:
                    Wv = filterbank.analyze(noisy)
                    Wv_st = pt.sign(Wv) * F.relu(Wv.abs() - lam_t)
                    denoised = noisy - filterbank.synthesize(Wv - Wv_st)
                    psnr_sum += -10 * math.log10(F.mse_loss(denoised, x).item() + 1e-12)
                psnrs[ii] = psnr_sum / len(data_val)

            best_ind = np.argmax(psnrs)
            best_lam = float(lams[best_ind])

            if plot_result:
                fig, axs = plt.subplots(1, 1)
                axs.plot(lams, psnrs)
                axs.set_xscale("log")
                axs.set_yscale("log")
                axs.set_ylabel("PSNR [dB]")
                axs.stem(best_lam, psnrs[best_ind], linefmt="C1-.")
                axs.grid()
                axs.set_xlim(lams[0], lams[-1])
                axs.set_ylim(psnrs.min() * 0.95, psnrs.max() * 1.05)
                fig.tight_layout()

            res = np.full(weights_shape, best_lam)
            if filter_weights is not None:
                res *= filter_weights.reshape(weights_shape)
            return res

        else:
            raise ValueError(f"Unknown option: {method}")

huber_moreau_prox 

huber_moreau_prox(
    u: Tensor, sigma: float, lam: Tensor, delta: Tensor
) -> Tensor

Proximal operator of sigma * (lamphi_delta)^ via Moreau decomposition:

prox_{sigma*f^*}(u) = u - sigma * prox_{f/sigma}(u/sigma)

Used in the PDHG dual update. Differentiable w.r.t. u, lam, delta.

Source code in src/autoden/transforms/learnable_regularizers.py 
148
149
150
151
152
153
154
155
156
157
def huber_moreau_prox(u: pt.Tensor, sigma: float, lam: pt.Tensor, delta: pt.Tensor) -> pt.Tensor:
    """
    Proximal operator of  sigma * (lam*phi_delta)^*  via Moreau decomposition:

        prox_{sigma*f^*}(u) = u - sigma * prox_{f/sigma}(u/sigma)

    Used in the PDHG dual update.
    Differentiable w.r.t. u, lam, delta.
    """
    return u - sigma * huber_prox(u / sigma, lam / sigma, delta)

huber_prox 

huber_prox(v: Tensor, lam: Tensor, delta: Tensor) -> Tensor

Proximal operator of lam * phi_delta applied element-wise.

prox_{lamphi_delta}(v) = delta/(delta+lam) * v if |v| <= delta + lam (scaled shrinkage) v - lamsign(v) if |v| > delta + lam (soft threshold)

All inputs broadcast freely. Differentiable w.r.t. v, lam, delta.

Source code in src/autoden/transforms/learnable_regularizers.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def huber_prox(v: pt.Tensor, lam: pt.Tensor, delta: pt.Tensor) -> pt.Tensor:
    """
    Proximal operator of  lam * phi_delta  applied element-wise.

    prox_{lam*phi_delta}(v) =
        delta/(delta+lam) * v      if |v| <= delta + lam   (scaled shrinkage)
        v - lam*sign(v)            if |v| >  delta + lam   (soft threshold)

    All inputs broadcast freely.
    Differentiable w.r.t. v, lam, delta.
    """
    v_abs = v.abs()
    thresh = delta + lam
    # Scaled shrinkage branch (quadratic region of Huber)
    shrunk = (delta / thresh) * v
    # Soft-threshold branch (linear region of Huber)
    soft = (v_abs - lam).clamp(min=0.0) * v.sign()
    return pt.where(v_abs <= thresh, shrunk, soft)

train_lambdas_denoising 

train_lambdas_denoising(
    regularizer: LearnableRegularizerL1,
    data_trn: NDArray,
    data_val: NDArray,
    sigma: float = 25.0 / 255.0,
    n_epochs: int = 50,
    batch_size: int = 16,
    lr: float = 0.001,
    sched_starts: int = 0,
    device: str = "cuda" if is_available() else "cpu",
    verbose: bool = True,
) -> tuple[LearnableRegularizerL1, NDArray]

Learn the per-filter thresholds/weights lambda_i by minimizing the MSE denoising loss:

L = E[||prox_R(x + epsilon) - x||^2]     epsilon ~ N(0, sigma^2 * I)

This decouples filter shape (learned by sparsity, task-agnostic) from threshold calibration (learned by denoising, noise-level specific).

You can re-run this phase for different sigma values without relearning filters.

Parameters:

  • regularizer (ParsevalL1Regularizer) –

    Must be a ConvolutionalDecompositionBase regularizer.

  • data_trn (NDArray) –

    Clean training images.

  • data_val (NDArray) –

    Clean validation images.

  • sigma (float, default: 25.0 / 255.0 ) –

    Noise standard deviation to calibrate for (default is 25.0 / 255.0).

  • n_epochs (int, default: 50 ) –

    Number of training epochs (default is 50).

  • batch_size (int, default: 16 ) –

    Batch size for training (default is 16).

  • lr (float, default: 0.001 ) –

    Learning rate (default is 1e-3).

  • sched_starts (int, default: 0 ) –

    Epoch at which the learning rate scheduler starts (default is 0).

  • device (str, default: 'cuda' if is_available() else 'cpu' ) –

    Device to use for training (default is "cuda" if available, else "cpu").

  • verbose (bool, default: True ) –

    Whether to print training progress (default is True).

Returns:

  • tuple[ParsevalL1Regularizer, NDArray] –

    A tuple containing the trained regularizer and the validation loss history.

Source code in src/autoden/transforms/learnable_regularizers.py 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
def train_lambdas_denoising(
    regularizer: LearnableRegularizerL1,
    data_trn: NDArray,  # clean images
    data_val: NDArray,  # clean images
    sigma: float = 25.0 / 255.0,
    n_epochs: int = 50,
    batch_size: int = 16,
    lr: float = 1e-3,
    sched_starts: int = 0,
    device: str = "cuda" if pt.cuda.is_available() else "cpu",
    verbose: bool = True,
) -> tuple[LearnableRegularizerL1, NDArray]:
    """
    Learn the per-filter thresholds/weights lambda_i by minimizing the MSE denoising loss:

        L = E[||prox_R(x + epsilon) - x||^2]     epsilon ~ N(0, sigma^2 * I)

    This decouples filter shape (learned by sparsity, task-agnostic) from
    threshold calibration (learned by denoising, noise-level specific).

    You can re-run this phase for different sigma values without relearning filters.

    Parameters
    ----------
    regularizer : ParsevalL1Regularizer
        Must be a ConvolutionalDecompositionBase regularizer.
    data_trn : NDArray
        Clean training images.
    data_val : NDArray
        Clean validation images.
    sigma : float, optional
        Noise standard deviation to calibrate for (default is 25.0 / 255.0).
    n_epochs : int, optional
        Number of training epochs (default is 50).
    batch_size : int, optional
        Batch size for training (default is 16).
    lr : float, optional
        Learning rate (default is 1e-3).
    sched_starts : int, optional
        Epoch at which the learning rate scheduler starts (default is 0).
    device : str, optional
        Device to use for training (default is "cuda" if available, else "cpu").
    verbose : bool, optional
        Whether to print training progress (default is True).

    Returns
    -------
    tuple[ParsevalL1Regularizer, NDArray]
        A tuple containing the trained regularizer and the validation loss history.
    """
    reg = regularizer.to(device)

    # Freeze the filterbank completely
    # reg.fb.A.requires_grad_(False)

    trn_dset = DatasetNumpy(data_trn, device)
    trn_dsets_list = DatasetsList([trn_dset, trn_dset], augmentation=["flip", "rot", AugmentationGaussianNoise(sigma)])

    val_dset = DatasetNumpy(data_val, device)
    val_dsets_list = DatasetsList([val_dset, val_dset], augmentation=["flip", "rot", AugmentationGaussianNoise(sigma)])

    trn_dl = DataLoader(trn_dsets_list, batch_size=batch_size, shuffle=True, num_workers=0)
    # , pin_memory=(device == "cuda")
    val_dl = DataLoader(val_dsets_list, batch_size=batch_size, shuffle=False, num_workers=0)

    opt = pt.optim.Adam(reg.trainable_params(), lr=lr)
    if sched_starts > 0:
        sch = pt.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=n_epochs)
    else:
        sch = None

    history = {"loss_trn": [], "loss_val": [], "lambda_mean": []}
    best_loss_val = float("inf")
    best_lambdas = reg.lambdas.detach().clone()

    if verbose:
        print(f"\nPhase 2 - Lambda calibration (denoising, sigma={sigma:.4f}) for  filter bank: k={reg.fb.k}, m={reg.fb.m}")
        print(f"  Filters: FROZEN  |  lambda_i: learning  |  fix_lambda0={reg.fix_dc}")

    for epoch in range(1, n_epochs + 1):
        reg.train()
        total_trn = 0.0
        for noisy, clean in trn_dl:
            denoised = reg.prox(noisy)
            loss = F.mse_loss(denoised, clean)
            opt.zero_grad()
            loss.backward()
            opt.step()
            total_trn += loss.item()
        if sch is not None:
            sch.step()
        loss_trn = total_trn / len(trn_dl)

        # ── validate ────────────────────────────────────────────────────────
        reg.eval()
        total_val = 0.0
        with pt.no_grad():
            for noisy, clean in val_dl:
                denoised = reg.prox(noisy)
                total_val += F.mse_loss(denoised, clean).item()
        loss_val = total_val / max(len(val_dl), 1)

        lam_mean = reg.lambdas.mean().item()

        history["loss_trn"].append(loss_trn)
        history["loss_val"].append(loss_val)
        history["lambda_mean"].append(lam_mean)

        if verbose and (epoch % 10 == 0 or epoch == 1):
            psnr = -10 * math.log10(loss_val + 1e-12)

            lams = reg.lambdas.squeeze()
            print(
                f"  epoch {epoch:4d}/{n_epochs}  "
                f"train={loss_trn:.5f}  val={loss_val:.5f}  val_PSNR={psnr:.2f}dB"
                f" - λ_0={lams[0]:.5f}  λ=[{lams[1:].min():.5f}, {lams[1:].max():.5f}]",
                end="",
                flush=True,
            )
            if isinstance(reg, LearnableRegularizerHuber):
                dels = reg.deltas.squeeze()
                print(f" - delta_0={dels[0]:.5f}  " f"deltas: min={dels[1:].min():.5f} max={dels[1:].max():.5f}")
            else:
                print("")

        if loss_val < best_loss_val:
            best_loss_val = loss_val
            best_lambdas = reg.lambdas.detach().clone()

    if verbose:
        print(f"  Best val loss: {best_loss_val:.6f}")
    return reg, best_lambdas.cpu().numpy().copy()