learning.ComplexNet

learning.ComplexNet#

class bartorch.learning.ComplexNet#

Bases: Module

Network on real channels applied to complex images whose leading axes are channels.

The input is (batch, *channels, *spatial), or (batch, *channels, frames, *spatial) with frames. The real and imaginary parts and the channels axes become the network’s 2 * prod(channels) real channels, real parts first; a frame axis stays in front of the spatial ones, as UNet with frames=True takes it. Subspace coefficient maps are thereby denoised jointly as one multi-channel image.

Parameters:

  • net (callable) – Called on real (batch, 2 * prod(channels), [frames,] *spatial) tensors as net(x) or net(x, sigma), with any further keyword this module is called with, such as step. Returns its input’s shape.

  • spatial ({2, 3}, default=3) – Number of trailing spatial axes.

  • channels (int, default=0) – Number of axes in front of the spatial ones (and of the frame axis) folded into the channels.

  • frames (bool, default=False) – Whether the axis in front of the spatial ones is a frame axis.

  • normalize ({"peak", "whiten", None}, default="peak") – Scaling of each item around the call, with statistics taken without gradient. "peak" divides by the peak modulus, so sigma is relative to it; "whiten" subtracts each real channel’s mean and decorrelates the channels to unit variance, balancing coefficient maps whose energies differ by orders of magnitude, so sigma has no fixed unit; None leaves the item as it is.

Raises:

  • TypeError – net is not callable.

  • ValueError – An unknown normalize.

Examples

>>> net = learning.ComplexNet(learning.UNet(10, steps=True), channels=1, normalize="whiten")
>>> prior = priors.ImplicitPrior(net, step=True)      # (5, z, y, x) coefficient maps
forward()#

Apply the network.

Parameters:

  • input (torch.Tensor) – (batch, *channels, [frames,] *spatial). A real tensor is taken as complex with a zero imaginary part.

  • sigma (float or torch.Tensor, default=None) – Passed to the network as its second argument when given.

  • **keywords – Passed to the network.

Returns:

Of input’s shape and dtype; the real part for a real input.

Return type:

torch.Tensor