priors.Regularizer#

class bartorch.priors.Regularizer#

Bases: ABC

Regularization term \(g(G x)\) whose proximal operator and transform BART builds.

\(G\) is a linear transform and \(g\) a functional on its codomain, both from BART’s opt_reg_configure. \(G\) is the identity for a term on the image itself and for a term whose transform is inside its proximal operator, such as Wavelet; it is a proper operator for a term such as TotalVariation, whose functional acts on finite differences. Only the alternating-direction and primal-dual iterations apply \(G\); see transform_is_identity().

Both are built per image shape and cached, so solving twice with the same term builds nothing the second time.

kind#

BART’s letter for the term, as -R takes it.

Type:

str

weight#

The term’s weight \(\lambda\), relative to data divided by data_scaling().

Type:

float

axes#

Axes the transform acts over, as indices into the image’s shape.

Type:

tuple of int

joint_axes#

Axes along which entries are thresholded jointly, as one \(\ell_2\) norm.

Type:

tuple of int

count#

For an NIHT term, the percentage of entries kept; zero otherwise.

Type:

int

Examples

>>> term = priors.Wavelet((-1, -2), 0.01)
>>> z = term.prox(x, 0.5)                    # BART's proximal operator, step 0.5
>>> term.transform_is_identity((64, 64))     # G = I: every proximal solver takes it
True
>>> priors.TotalVariation((-1, -2), 0.01).transform_is_identity((64, 64))
False
build()#

Handle of BART’s operator for this term over an image of C-order shape.

Built on first use for each shape and cached. A term that draws random shifts is built once per batch item, so that each item draws the sequence it would draw alone. The handle is owned by the term and freed with it.

Parameters:
  • shape (tuple of int) – Image shape.

  • item (int, default=0) – Batch item.

Returns:

The handle.

Return type:

int

Raises:

BartError – BART could not build the term over shape.

prox_shape()#

C-order shape of the domain of the proximal operator, for an image of image_shape.

The image’s shape when \(G\) is the identity, wavelet terms included; otherwise the codomain of \(G\), which for total variation carries the gradient components on an axis of their own in front.

Parameters:

image_shape (tuple of int) – Image shape.

Return type:

tuple of int

prox()#

Apply the proximal operator \(\operatorname{prox}_{\gamma g}\) as BART’s solvers do.

Parameters:
  • x (tensor) – Of prox_shape(), the image’s shape for most terms.

  • gamma (float, default=1.0) – Step \(\gamma\); the term’s weight is already in \(g\).

  • image_shape (tuple of int, default=None) – Image shape the term is configured for, when it differs from x.shape, as for total variation; x.shape when omitted.

  • item (int, default=0) – Batch item x is; a term that draws random shifts keeps a generator per item.

Returns:

Complex64, of x’s shape.

Return type:

torch.Tensor

Raises:
  • ValueError – x requires a gradient: BART’s proximal operators have no backward pass. Use ImplicitPrior for a differentiable proximal step, or detach() to hold this one fixed. Also raised for an x not of prox_shape().

  • BartError – BART could not apply the operator.

Notes

This is the proximal operator alone; a solver applying the term computes it on \(G x\) (apply_transform()). BART’s IST and FISTA apply it to the image and ignore \(G\), so they take only terms whose transform is the identity.

Examples

>>> priors.Wavelet((-1, -2), 0.01).prox(image, gamma=0.95)
apply_transform()#

Apply \(G\), its adjoint or \(G^H G\) to x.

Works at any rank, including total variation’s, whose codomain exceeds BART’s 16 dimensions and so has no transform() operator.

Parameters:
  • x (tensor) – The image for "forward" and "normal", of prox_shape() for "adjoint".

  • image_shape (tuple of int, default=None) – Image shape the term is configured for; x.shape when omitted, which holds for "forward" and "normal".

  • mode ({"forward", "adjoint", "normal"}, default="forward") – \(G\), \(G^H\) or \(G^H G\).

Returns:

Complex64.

Return type:

torch.Tensor

Raises:

ValueError – An unknown mode, or x not of the shape mode takes.

Notes

Recorded for autograd when x requires a gradient: the backward pass of "forward" is "adjoint" and vice versa, and "normal" is its own.

transform()#

The transform \(G\) as a LinearOperator.

The identity for a term whose transform is inside its proximal operator. The shapes do not tell the two arrangements apart: the Laplace term’s transform is a convolution whose codomain has the image’s shape.

Parameters:

image_shape (tuple of int) – Image shape.

Returns:

\(G\), from the image to prox_shape().

Return type:

LinearOperator

Raises:

BartError – The transform’s rank exceeds BART’s 16 dimensions, as for total variation; apply_transform() applies it instead.

rewind()#

Reset the term’s own random generator, for an image of image_shape.

A wavelet or locally low-rank term shifts its transform by a draw from a generator seeded when BART builds the operator. pics builds a fresh operator per run; every solver here rewinds its terms before it iterates, so a reused term draws as the command does. A term without a generator is unaffected.

Parameters:

image_shape (tuple of int) – Image shape.

transform_is_identity()#

Whether \(G\) is the identity, as BART’s linop_is_identity decides.

Decided on the operator, not from shapes: the Laplace term’s transform is a convolution on the image’s own shape. iter2_chambolle_pock asks this of the first term to choose it as the primal proximal step.

Parameters:

image_shape (tuple of int) – Image shape.

Returns:

True when \(G\) is the identity.

Return type:

bool

detach()#

Return this term with its proximal step detached from the autograd graph.

The detached term thresholds exactly as before, and the gradient obtained is the iteration’s with this term held fixed: the use is a solve in which a differentiable denoiser occupies one term and a BART term another, and only the denoiser is trained.

Returns:

A view of this term whose prox() detaches its input.

Return type:

Regularizer

Examples

>>> optim.admm(y, A, [denoiser, priors.Wavelet((-1, -2), 0.01).detach()])