nlop.NonlinearOperator

nlop.NonlinearOperator#

class bartorch.nlop.NonlinearOperator#

Bases: Operator

Nonlinear operator on complex tensors, with its derivative and the derivative’s adjoint.

F(x) applies the operator; when an input requires a gradient the application is recorded for autograd, and the backward pass applies the adjoint of the derivative at x. a @ b composes into one BART operator, applying b first; either side may be a LinearOperator. partial() fixes one input, and linearize() returns the derivative at a point as a linear operator.

An operator may have several inputs and several outputs (BART’s nlop_s). Operators are defined in Python by from_callbacks() or TorchOperator.

ishapes, oshapes

Shape of each input and of each output, C order.

Type:

tuple of tuple of int

ishape, oshape

Domain and codomain of an operator with one input and one output.

Type:

tuple of int

Examples

>>> M = nlop.MultiEcho([10.0, 20.0, 40.0], (64, 64))    # echo times in ms
>>> theta = M.initial(T2=80.0)
>>> M(theta).shape
torch.Size([3, 64, 64])
>>> D = M.linearize(theta)            # the derivative at theta, a LinearOperator
>>> E = linop.FFT(M.oshape, axes=(-1, -2)) @ M    # still a NonlinearOperator
property ishape#

Domain of an operator with one input; ValueError for several.

property oshape#

Codomain of an operator with one output; ValueError for several.

forward()#

Evaluate F(*xs) without recording it for autograd.

The evaluation point becomes the point at which the operator’s derivative and adjoint are taken until the next evaluation.

Parameters:

*xs (torch.Tensor) – One tensor per input, of the shapes in ishapes.

Returns:

One complex64 tensor per output; the tensor itself for one output.

Return type:

torch.Tensor or tuple of torch.Tensor

Raises:

ValueError – If the number of tensors is not the number of inputs, or the inputs are on different devices.

classmethod from_callbacks()#

Build an operator from Python functions, applied through BART’s callbacks.

Parameters:

  • oshape (tuple of int or sequence of tuple of int) – Codomain and domain, C order. A sequence of shapes gives that many outputs or inputs.

  • ishape (tuple of int or sequence of tuple of int) – Codomain and domain, C order. A sequence of shapes gives that many outputs or inputs.

  • forward (callable) – forward(*xs), one tensor per input, returning one tensor per output. Each call sets the point at which derivative and adjoint are taken.

  • derivative (callable) – derivative(dx) and adjoint(dy) for one input and one output; derivative(o, i, dx) and adjoint(o, i, dy), the derivative of output o by input i and its adjoint, for several.

  • adjoint (callable) – derivative(dx) and adjoint(dy) for one input and one output; derivative(o, i, dx) and adjoint(o, i, dy), the derivative of output o by input i and its adjoint, for several.

Return type:

NonlinearOperator

linearize()#

Return the derivative at a point as a LinearOperator.

Parameters:

  • *xs (torch.Tensor) – The point, one tensor per input.

  • input (int, default=None) – Input to differentiate by, the others held at their values. None differentiates by all inputs, flattened and concatenated in input order into one vector, as a Gauss-Newton step does.

  • output (int, default=0) – Output to differentiate.

Returns:

DF(x). It holds the point: later evaluations of the operator do not change it, and it is differentiable by x. For an operator that does not supply its derivative as a function of the point, it is instead taken at the last evaluated point and follows later evaluations.

Return type:

LinearOperator

Raises:

  • ValueError – If the number of tensors is not the number of inputs.

  • IndexError – If input or output is out of range.

__call__()#

Evaluate F(*xs), recorded for autograd when an input requires a gradient.

Parameters:

*xs (torch.Tensor) – One tensor per input.

Returns:

One tensor per output; the tensor itself for one output.

Return type:

torch.Tensor or tuple of torch.Tensor

partial()#

Fix one input to a constant, BART’s nlop_set_input_const.

Parameters:

  • input (int) – Input to fix; negative counts from the last.

  • value (torch.Tensor) – Its value, of that input’s shape.

Returns:

The operator without that input; the remaining inputs keep their order.

Return type:

NonlinearOperator