scico.linop

Linear operator functions and classes.

Modules

scico.linop.optics

Optical propagator classes.

scico.linop.xray

X-ray transform classes.

Functions

jacobian(F, u[, include_eval])

Construct Jacobian linear operator for a general operator.

linop_from_function(f, classname[, f_name])

Make a LinearOperator from a function.

linop_over_axes(linop, input_shape, *args[, ...])

Construct a list of LinearOperator by iterating over axes.

operator_norm(A[, maxiter, key])

Estimate the norm of a LinearOperator.

power_iteration(A[, maxiter, key])

Compute largest eigenvalue of a diagonalizable LinearOperator.

valid_adjoint(A, AT[, eps, x, y, key])

Check whether LinearOperator AT is the adjoint of A.

Classes

CircularConvolve(h, input_shape[, ndims, ...])

A circular convolution linear operator.

ComposedLinearOperator(A, B[, jit])

A composition of two LinearOperator objects.

Convolve(h, input_shape[, input_dtype, ...])

A convolution linear operator.

Crop(crop_width, input_shape[, input_dtype, jit])

A linear operator for cropping an array.

CylindricalGradient(input_shape[, axes, ...])

Gradient projected into cylindrical coordinates.

DFT(input_shape[, axes, axes_shape, norm, jit])

Multi-dimensional discrete Fourier transform.

Diagonal(diagonal[, input_shape, input_dtype])

Diagonal linear operator.

DiagonalReplicated(op, replicates[, ...])

A diagonal stack constructed from a single linear operator.

DiagonalStack(ops[, collapse_input, ...])

A diagonal stack of linear operators.

FiniteDifference(input_shape[, input_dtype, ...])

Finite difference operator.

Identity(input_shape[, input_dtype])

Identity operator.

LinearOperator(input_shape[, output_shape, ...])

Generic linear operator base class

MatrixOperator(A[, input_cols])

Linear operator implementing matrix multiplication.

Pad(input_shape, *args[, input_dtype, ...])

Linear operator version of scico.numpy.pad.

PolarGradient(input_shape[, axes, center, ...])

Gradient projected into polar coordinates.

ProjectedGradient(input_shape[, axes, ...])

Gradient projected onto local coordinate system.

Reshape(input_shape, *args[, input_dtype, ...])

Linear operator version of scico.numpy.reshape.

ScaledIdentity(scalar, input_shape[, ...])

Scaled identity operator.

SingleAxisFiniteDifference(input_shape[, ...])

Finite difference operator acting along a single axis.

Slice(idx, input_shape[, input_dtype, jit])

A linear operator for slicing an array.

SphericalGradient(input_shape[, axes, ...])

Gradient projected into spherical coordinates.

Sum(input_shape, *args[, input_dtype, ...])

Linear operator version of scico.numpy.sum.

Transpose(input_shape, *args[, input_dtype, ...])

Linear operator version of scico.numpy.transpose.

VerticalStack(ops[, collapse_output, jit])

A vertical stack of linear operators.
class scico.linop.CircularConvolve(h, input_shape, ndims=None, input_dtype=<class 'jax.numpy.float32'>, h_is_dft=False, h_center=None, jit=True, **kwargs)

Bases: LinearOperator

A circular convolution linear operator.

Inheritance diagram of CircularConvolve

This linear operator implements circular, multi-dimensional convolution via pointwise multiplication in the DFT domain. In its simplest form, it implements a single convolution and can be represented by linear operator \(H\) such that

\[H \mb{x} = \mb{h} \ast \mb{x} \;,\]

where \(\mb{h}\) is a user-defined filter.

More complex forms, corresponding to the case where either the input (as represented by parameter input_shape) or filter (parameter h) have additional axes that are not involved in the convolution are also supported. These follow numpy broadcasting rules. For example:

Additional axes in the input \(\mb{x}\) and not in \(\mb{h}\) corresponds to the operation

\[\begin{split}H \mb{x} = \left( \begin{array}{ccc} H' & 0 & \ldots\\ 0 & H' & \ldots\\ \vdots & \vdots & \ddots \end{array} \right) \left( \begin{array}{c} \mb{x}_0\\ \mb{x}_1\\ \vdots \end{array} \right) \;.\end{split}\]

Additional axes in \(\mb{h}\) corresponds to multiple filters, which will be denoted by \(\{\mb{h}_m\}\), with corresponding individual linear operations being denoted by \(h_m \mb{x}_m = \mb{h}_m \ast \mb{x}_m\). The full linear operator can then be represented as

\[\begin{split}H \mb{x} = \left( \begin{array}{c} H_0\\ H_1\\ \vdots \end{array} \right) \mb{x} \;.\end{split}\]

if the input is singleton, and as

\[\begin{split}H \mb{x} = \left( \begin{array}{ccc} H_0 & 0 & \ldots\\ 0 & H_1 & \ldots\\ \vdots & \vdots & \ddots \end{array} \right) \left( \begin{array}{c} \mb{x}_0\\ \mb{x}_1\\ \vdots \end{array} \right)\end{split}\]

otherwise.

Parameters:

  • h (Array) – Array of filters.

  • input_shape (Tuple[int, ...]) – Shape of input array.

  • ndims (Optional[int]) – Number of (trailing) dimensions of the input and h involved in the convolution. Defaults to the number of dimensions in the input.

  • input_dtype (DType) – dtype for input argument. Defaults to float32.

  • h_is_dft (bool) – Flag indicating whether h is in the DFT domain.

  • h_center (Union[Array, ndarray, Sequence, float, int, None]) – Array of length ndims specifying the center of the filter. Defaults to the upper left corner, i.e., h_center = [0, 0, …, 0], may be noninteger. May be a float or int if h is one-dimensional.

  • jit (bool) – If True, jit the evaluation, adjoint, and gram functions of the LinearOperator.

static from_operator(H, ndims=None, center=None, jit=True)[source]

Construct a CircularConvolve version of a given operator.

Construct a CircularConvolve version of a given operator, which is assumed to be linear and shift invariant (LSI).

Parameters:

  • H (Operator) – Input operator.

  • ndims (Optional[int]) – Number of trailing dims over which the H acts.

  • center (Optional[Tuple[int, ...]]) – Location at which to place the Kronecker delta. For LSI inputs, this will not matter. Defaults to the center of H.input_shape, i.e., (n_1 // 2, n_2 // 2, …).

  • jit (bool) – If True, jit the resulting CircularConvolve.

class scico.linop.Convolve(h, input_shape, input_dtype=<class 'numpy.float32'>, mode='full', jit=True, **kwargs)

Bases: LinearOperator

A convolution linear operator.

Inheritance diagram of Convolve

Wrap jax.scipy.signal.convolve as a LinearOperator.

Parameters:

  • h (Array) – Convolutional filter. Must have same number of dimensions as len(input_shape).

  • input_shape (Tuple[int, ...]) – Shape of input array.

  • input_dtype (DType) – dtype for input argument. Defaults to float32.

  • mode (str) – A string indicating the size of the output. One of “full”, “valid”, “same”. Defaults to “full”.

  • jit (bool) – If True, jit the evaluation, adjoint, and gram functions of the LinearOperator.

For more details on mode, see jax.scipy.signal.convolve.

class scico.linop.DFT(input_shape, axes=None, axes_shape=None, norm=None, jit=True, **kwargs)

Bases: LinearOperator

Multi-dimensional discrete Fourier transform.

Inheritance diagram of DFT

Parameters:

  • input_shape (Tuple[int, ...]) – Shape of input array.

  • axes (Optional[Sequence]) – Axes over which to compute the DFT. If None, the DFT is computed over all axes.

  • axes_shape (Optional[Tuple[int, ...]]) – Output shape on the subset of array axes selected by axes. This parameter has the same behavior as the s parameter of numpy.fft.fftn.

  • norm (Optional[str]) – DFT normalization mode. See the norm parameter of numpy.fft.fftn.

  • jit (bool) – If True, jit the evaluation, adjoint, and gram functions of the LinearOperator.

inv(z)[source]

Compute the inverse of this LinearOperator.

Compute the inverse of this LinearOperator applied to z.

Parameters:

z (Array) – Input array to inverse DFT.

Return type:

Array

class scico.linop.Diagonal(diagonal, input_shape=None, input_dtype=None, **kwargs)

Bases: LinearOperator

Diagonal linear operator.

Inheritance diagram of Diagonal

Parameters:

  • diagonal (Union[Array, BlockArray]) – Diagonal elements of this LinearOperator.

  • input_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...], None]) – Shape of input array. By default, equal to diagonal.shape, but may also be set to a shape that is broadcast-compatible with diagonal.shape.

  • input_dtype (Optional[DType]) – dtype of input argument. The default, None, means diagonal.dtype.

property H: Diagonal

Hermitian transpose of this Diagonal.

property T: Diagonal

Transpose of this Diagonal.

conj()[source]

Complex conjugate of this Diagonal.

Return type:

Diagonal

property diagonal: Array | BlockArray

Return an array representing the diagonal component.

property gram_op: Diagonal

Gram operator of this Diagonal.

Return a new Diagonal G such that G(x) = A.adj(A(x))).

norm(ord=None)[source]

Compute the matrix norm of the diagonal operator.

Valid values of ord and the corresponding norm definition are those listed under “norm for matrices” in the scico.numpy.linalg.norm documentation.

class scico.linop.Identity(input_shape, input_dtype=<class 'jax.numpy.float32'>, **kwargs)

Bases: ScaledIdentity

Identity operator.

Inheritance diagram of Identity

Parameters:
conj()[source]

Complex conjugate of this Identity.

Return type:

Identity

property gram_op: Identity

Gram operator of this Identity.

class scico.linop.ScaledIdentity(scalar, input_shape, input_dtype=<class 'jax.numpy.float32'>, **kwargs)

Bases: Diagonal

Scaled identity operator.

Inheritance diagram of ScaledIdentity

Parameters:
conj()[source]

Complex conjugate of this ScaledIdentity.

Return type:

ScaledIdentity

property diagonal: Array | BlockArray

Return an array representing the diagonal component.

property gram_op: ScaledIdentity

Gram operator of this ScaledIdentity.

norm(ord=None)[source]

Compute the matrix norm of the identity operator.

Valid values of ord and the corresponding norm definition are those listed under “norm for matrices” in the scico.numpy.linalg.norm documentation.

class scico.linop.FiniteDifference(input_shape, input_dtype=<class 'numpy.float32'>, axes=None, prepend=None, append=None, circular=False, jit=True, **kwargs)

Bases: VerticalStack

Finite difference operator.

Inheritance diagram of FiniteDifference

Compute finite differences along the specified axes, returning the results in a jax.Array (when possible) or BlockArray. See VerticalStack for details on how this choice is made. See SingleAxisFiniteDifference for the mathematical implications of the different boundary handling options prepend, append, and circular.

Example

>>> A = FiniteDifference((2, 3))
>>> x = snp.array([[1, 2, 4],
...                [0, 4, 1]])
>>> (A @ x)[0]
Array([[-1,  2, -3]], dtype=int32)
>>> (A @ x)[1]
Array([[ 1,  2],
       [ 4, -3]], dtype=int32)
Parameters:

  • input_shape (Tuple[int, ...]) – Shape of input array.

  • input_dtype (DType) – dtype for input argument. Defaults to float32.

  • axes (Union[int, Tuple[int, ...], List[int], None]) – Axis or axes over which to apply finite difference operator. If not specified, or None, differences are evaluated along all axes.

  • prepend (Union[Literal[0], Literal[1], None]) – Flag indicating handling of the left/top/etc. boundary. If None, there is no boundary extension. Values of 0 or 1 indicate respectively that zeros or the initial value in the array are prepended to the difference array.

  • append (Union[Literal[0], Literal[1], None]) – Flag indicating handling of the right/bottom/etc. boundary. If None, there is no boundary extension. Values of 0 or 1 indicate respectively that zeros or -1 times the final value in the array are appended to the difference array.

  • circular (bool) – If True, perform circular differences, i.e., include x[-1] - x[0]. If True, prepend and append must both be None.

  • jit (bool) – If True, jit the evaluation, adjoint, and gram functions of the LinearOperator.

class scico.linop.SingleAxisFiniteDifference(input_shape, input_dtype=<class 'numpy.float32'>, axis=-1, prepend=None, append=None, circular=False, jit=True, **kwargs)

Bases: LinearOperator

Finite difference operator acting along a single axis.

Inheritance diagram of SingleAxisFiniteDifference

By default (i.e. prepend and append set to None and circular set to False), the difference operator corresponds to the matrix

\[\begin{split}\left(\begin{array}{rrrrr} -1 & 1 & 0 & \ldots & 0\\ 0 & -1 & 1 & \ldots & 0\\ \vdots & \vdots & \ddots & \ddots & \vdots\\ 0 & 0 & \ldots & -1 & 1 \end{array}\right) \;,\end{split}\]

mapping \(\mbb{R}^N \rightarrow \mbb{R}^{N-1}\), while if circular is True, it corresponds to the \(\mbb{R}^N \rightarrow \mbb{R}^N\) mapping

\[\begin{split}\left(\begin{array}{rrrrr} -1 & 1 & 0 & \ldots & 0\\ 0 & -1 & 1 & \ldots & 0\\ \vdots & \vdots & \ddots & \ddots & \vdots\\ 0 & 0 & \ldots & -1 & 1\\ 1 & 0 & \dots & 0 & -1 \end{array}\right) \;.\end{split}\]

Other possible choices include prepend set to None and append set to 0, giving the \(\mbb{R}^N \rightarrow \mbb{R}^N\) mapping

\[\begin{split}\left(\begin{array}{rrrrr} -1 & 1 & 0 & \ldots & 0\\ 0 & -1 & 1 & \ldots & 0\\ \vdots & \vdots & \ddots & \ddots & \vdots\\ 0 & 0 & \ldots & -1 & 1\\ 0 & 0 & \dots & 0 & 0 \end{array}\right) \;,\end{split}\]

and both prepend and append set to 1, giving the \(\mbb{R}^N \rightarrow \mbb{R}^{N+1}\) mapping

\[\begin{split}\left(\begin{array}{rrrrr} 1 & 0 & 0 & \ldots & 0\\ -1 & 1 & 0 & \ldots & 0\\ 0 & -1 & 1 & \ldots & 0\\ \vdots & \vdots & \ddots & \ddots & \vdots\\ 0 & 0 & \ldots & -1 & 1\\ 0 & 0 & \dots & 0 & -1 \end{array}\right) \;.\end{split}\]
Parameters:

  • input_shape (Tuple[int, ...]) – Shape of input array.

  • input_dtype (DType) – dtype for input argument. Defaults to float32.

  • axis (int) – Axis over which to apply finite difference operator.

  • prepend (Union[Literal[0], Literal[1], None]) – Flag indicating handling of the left/top/etc. boundary. If None, there is no boundary extension. Values of 0 or 1 indicate respectively that zeros or the initial value in the array are prepended to the difference array.

  • append (Union[Literal[0], Literal[1], None]) – Flag indicating handling of the right/bottom/etc. boundary. If None, there is no boundary extension. Values of 0 or 1 indicate respectively that zeros or -1 times the final value in the array are appended to the difference array.

  • circular (bool) – If True, perform circular differences, i.e., include x[-1] - x[0]. If True, prepend and append must both be None.

  • jit (bool) – If True, jit the evaluation, adjoint, and gram functions of the LinearOperator.

class scico.linop.Crop(crop_width, input_shape, input_dtype=<class 'jax.numpy.float32'>, jit=True, **kwargs)

Bases: LinearOperator

A linear operator for cropping an array.

Inheritance diagram of Crop

Parameters:

  • crop_width (Union[int, Sequence]) – Specify the crop width using the same format as the pad_width parameter of snp.pad.

  • input_shape (Tuple[int, ...]) – Shape of input jax.Array.

  • input_dtype (DType) – dtype for input argument. Defaults to float32.

  • jit (bool) – If True, jit the evaluation, adjoint, and gram functions of the LinearOperator.

class scico.linop.Pad(input_shape, *args, input_dtype=<class 'jax.numpy.float32'>, output_shape=None, output_dtype=None, jit=True, **kwargs)

Bases: LinearOperator

Linear operator version of scico.numpy.pad.

Inheritance diagram of Pad

Parameters:

  • input_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]) – Shape of input array.

  • args (Any) – Positional arguments passed to scico.numpy.pad.

  • input_dtype (DType) – dtype for input argument. Defaults to float32. If the LinearOperator implements complex-valued operations, this must be a complex dtype (typically complex64) for correct adjoint and gradient calculation.

  • output_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...], None]) – Shape of output array. Defaults to None. If None, output_shape is determined by evaluating self.__call__ on an input array of zeros.

  • output_dtype (Optional[DType]) – dtype for output argument. Defaults to None. If None, output_dtype is determined by evaluating self.__call__ on an input array of zeros.

  • jit (bool) – If True, call jit on this LinearOperator to jit the forward, adjoint, and gram functions. Same as calling jit after the LinearOperator is created.

  • kwargs (Any) – Keyword arguments passed to scico.numpy.pad.

class scico.linop.Reshape(input_shape, *args, input_dtype=<class 'jax.numpy.float32'>, output_shape=None, output_dtype=None, jit=True, **kwargs)

Bases: LinearOperator

Linear operator version of scico.numpy.reshape.

Inheritance diagram of Reshape

Parameters:

  • input_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]) – Shape of input array.

  • args (Any) – Positional arguments passed to scico.numpy.reshape.

  • input_dtype (DType) – dtype for input argument. Defaults to float32. If the LinearOperator implements complex-valued operations, this must be a complex dtype (typically complex64) for correct adjoint and gradient calculation.

  • output_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...], None]) – Shape of output array. Defaults to None. If None, output_shape is determined by evaluating self.__call__ on an input array of zeros.

  • output_dtype (Optional[DType]) – dtype for output argument. Defaults to None. If None, output_dtype is determined by evaluating self.__call__ on an input array of zeros.

  • jit (bool) – If True, call jit on this LinearOperator to jit the forward, adjoint, and gram functions. Same as calling jit after the LinearOperator is created.

  • kwargs (Any) – Keyword arguments passed to scico.numpy.reshape.

class scico.linop.Slice(idx, input_shape, input_dtype=<class 'jax.numpy.float32'>, jit=True, **kwargs)

Bases: LinearOperator

A linear operator for slicing an array.

Inheritance diagram of Slice

This operator may be applied to either a jax.Array or a BlockArray. In the latter case, parameter idx must conform to the BlockArray indexing requirements.

Parameters:
class scico.linop.Sum(input_shape, *args, input_dtype=<class 'jax.numpy.float32'>, output_shape=None, output_dtype=None, jit=True, **kwargs)

Bases: LinearOperator

Linear operator version of scico.numpy.sum.

Inheritance diagram of Sum

Parameters:

  • input_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]) – Shape of input array.

  • args (Any) – Positional arguments passed to scico.numpy.sum.

  • input_dtype (DType) – dtype for input argument. Defaults to float32. If the LinearOperator implements complex-valued operations, this must be a complex dtype (typically complex64) for correct adjoint and gradient calculation.

  • output_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...], None]) – Shape of output array. Defaults to None. If None, output_shape is determined by evaluating self.__call__ on an input array of zeros.

  • output_dtype (Optional[DType]) – dtype for output argument. Defaults to None. If None, output_dtype is determined by evaluating self.__call__ on an input array of zeros.

  • jit (bool) – If True, call jit on this LinearOperator to jit the forward, adjoint, and gram functions. Same as calling jit after the LinearOperator is created.

  • kwargs (Any) – Keyword arguments passed to scico.numpy.sum.

class scico.linop.Transpose(input_shape, *args, input_dtype=<class 'jax.numpy.float32'>, output_shape=None, output_dtype=None, jit=True, **kwargs)

Bases: LinearOperator

Linear operator version of scico.numpy.transpose.

Inheritance diagram of Transpose

Parameters:

  • input_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]) – Shape of input array.

  • args (Any) – Positional arguments passed to scico.numpy.transpose.

  • input_dtype (DType) – dtype for input argument. Defaults to float32. If the LinearOperator implements complex-valued operations, this must be a complex dtype (typically complex64) for correct adjoint and gradient calculation.

  • output_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...], None]) – Shape of output array. Defaults to None. If None, output_shape is determined by evaluating self.__call__ on an input array of zeros.

  • output_dtype (Optional[DType]) – dtype for output argument. Defaults to None. If None, output_dtype is determined by evaluating self.__call__ on an input array of zeros.

  • jit (bool) – If True, call jit on this LinearOperator to jit the forward, adjoint, and gram functions. Same as calling jit after the LinearOperator is created.

  • kwargs (Any) – Keyword arguments passed to scico.numpy.transpose.

scico.linop.linop_from_function(f, classname, f_name=None)

Make a LinearOperator from a function.

Example

>>> Sum = linop_from_function(snp.sum, 'Sum')
>>> H = Sum((2, 10), axis=1)
>>> H @ snp.ones((2, 10))
Array([10., 10.], dtype=float32)
Parameters:

  • f (Callable) – Function from which to create a LinearOperator.

  • classname (str) – Name of the resulting class.

  • f_name (Optional[str]) – Name of f for use in docstrings. Useful for getting the correct version of wrapped functions. Defaults to f”{f.__module__}.{f.__name__}”.

class scico.linop.CylindricalGradient(input_shape, axes=None, center=None, angular=True, radial=True, axial=True, cdiff=False, input_dtype=<class 'numpy.float32'>, jit=True)

Bases: ProjectedGradient

Gradient projected into cylindrical coordinates.

Inheritance diagram of CylindricalGradient

Compute gradients projected onto cylindrical coordinate axes, as described in [29]. The local coordinate axes are illustrated in the figure below.

(png, hires.png, pdf)

../_images/cylindgrad.png

If only one of angular, radial, and axial is True, the operator output has the same shape as the input, otherwise the gradients for the selected local coordinate axes are stacked on an additional axis at index 0.

Parameters:

  • input_shape (Tuple[int, ...]) – Shape of input array.

  • axes (Optional[Tuple[int, ...]]) – Axes over which to compute the gradient. Should be a tuple \((i_x, i_y, i_z)\), where \(i_x\), \(i_y\) and \(i_z\) are input array axes assigned to \(x\), \(y\), and \(z\) coordinates respectively. Defaults to None, in which case the axes are taken to be (0, 1, 2). If an integer, this operator returns a jax.Array. If a tuple or None, the resulting arrays are stacked into a BlockArray.

  • center (Union[Tuple[int, ...], Array, None]) – Center of the cylindrical coordinate system in array indexing coordinates. Default is None, which places the center at the center of the two polar axes of the input array and at the zero index of the axial axis.

  • angular (bool) – Flag indicating whether to compute gradients in the angular (i.e. tangent to circles) direction.

  • radial (bool) – Flag indicating whether to compute gradients in the radial (i.e. directed outwards from the origin) direction.

  • axial (bool) – Flag indicating whether to compute gradients in the direction of the axis of the cylinder.

  • cdiff (bool) – If True, estimate gradients using the second order central different returned by snp.gradient, otherwise use the first order asymmetric difference returned by snp.diff.

  • input_dtype (DType) – dtype for input argument. Default is float32.

  • jit (bool) – If True, jit the evaluation, adjoint, and gram functions of the LinearOperator.

class scico.linop.PolarGradient(input_shape, axes=None, center=None, angular=True, radial=True, cdiff=False, input_dtype=<class 'numpy.float32'>, jit=True)

Bases: ProjectedGradient

Gradient projected into polar coordinates.

Inheritance diagram of PolarGradient

Compute gradients projected onto angular and/or radial axis directions, as described in [29]. Local coordinate axes are illustrated in the figure below.

(png, hires.png, pdf)

../_images/polargrad.png

If only one of angular and radial is True, the operator output has the same shape as the input, otherwise the gradients for the two local coordinate axes are stacked on an additional axis at index 0.

Parameters:

  • input_shape (Tuple[int, ...]) – Shape of input array.

  • axes (Optional[Tuple[int, ...]]) – Axes over which to compute the gradient. Should be a tuple \((i_x, i_y)\), where \(i_x\) and \(i_y\) are input array axes assigned to \(x\) and \(y\) coordinates respectively. Defaults to None, in which case the axes are taken to be (0, 1).

  • center (Union[Tuple[int, ...], Array, None]) – Center of the polar coordinate system in array indexing coordinates. Default is None, which places the center at the center of the input array.

  • angular (bool) – Flag indicating whether to compute gradients in the angular (i.e. tangent to circles) direction.

  • radial (bool) – Flag indicating whether to compute gradients in the radial (i.e. directed outwards from the origin) direction.

  • cdiff (bool) – If True, estimate gradients using the second order central different returned by snp.gradient, otherwise use the first order asymmetric difference returned by snp.diff.

  • input_dtype (DType) – dtype for input argument. Default is float32.

  • jit (bool) – If True, jit the evaluation, adjoint, and gram functions of the LinearOperator.

class scico.linop.ProjectedGradient(input_shape, axes=None, coord=None, cdiff=False, input_dtype=<class 'numpy.float32'>, jit=True)

Bases: LinearOperator

Gradient projected onto local coordinate system.

Inheritance diagram of ProjectedGradient

This class represents a linear operator that computes gradients of arrays projected onto a local coordinate system that may differ at every position in the array, as described in [29]. In the 2D illustration below \(x\) and \(y\) represent the standard coordinate system defined by the array axes, \((g_x, g_y)\) is the gradient vector within that coordinate system, \(x'\) and \(y'\) are the local coordinate axes, and \((g_x', g_y')\) is the gradient vector within the local coordinate system.

Figure illustrating projection of gradient onto local coordinate system.

Each of the local coordinate axes (e.g. \(x'\) and \(y'\) in the illustration above) is represented by a separate array in the coord tuple of arrays parameter of the class initializer.

Note

This operator should not be confused with the Projected Gradient optimization algorithm (a special case of Proximal Gradient), with which it is unrelated.

The result of applying the operator is always a jax.Array. If coord is a singleton tuple, it has the same shape as the input array. Otherwise, the gradients for each of the local coordinate axes are stacked on an additional axis at index 0.

If coord is None, which is the default, gradients are computed in the standard axis-aligned coordinate system, and the shape of the returned array depends on the number of axes on which the gradient is calculated, as specified explicitly or implicitly via the axes parameter.

Parameters:

  • input_shape (Tuple[int, ...]) – Shape of input array.

  • axes (Optional[Tuple[int, ...]]) – Axes over which to compute the gradient. Defaults to None, in which case the gradient is computed along all axes.

  • coord (Optional[Sequence[Union[Array, BlockArray]]]) – A tuple of arrays, each of which specifies a local coordinate axis direction. Each member of the tuple should either be a jax.Array or a BlockArray. If it is the former, it should have shape \(N \times M_0 \times M_1 \times \ldots\), where \(N\) is the number of axes specified by parameter axes, and \(M_i\) is the size of the \(i^{\mrm{th}}\) axis. If it is the latter, it should consist of \(N\) blocks, each of which has a shape that is suitable for multiplication with an array of shape \(M_0 \times M_1 \times \ldots\).

  • cdiff (bool) – If True, estimate gradients using the second order central different returned by snp.gradient, otherwise use the first order asymmetric difference returned by snp.diff.

  • input_dtype (DType) – dtype for input argument. Default is float32.

  • jit (bool) – If True, jit the evaluation, adjoint, and gram functions of the LinearOperator.

class scico.linop.SphericalGradient(input_shape, axes=None, center=None, azimuthal=True, polar=True, radial=True, cdiff=False, input_dtype=<class 'numpy.float32'>, jit=True)

Bases: ProjectedGradient

Gradient projected into spherical coordinates.

Inheritance diagram of SphericalGradient

Compute gradients projected onto spherical coordinate axes, based on the approach described in [29]. The local coordinate axes are illustrated in the figure below.

(png, hires.png, pdf)

../_images/spheregrad.png

If only one of azimuthal, polar, and radial is True, the operator output has the same shape as the input, otherwise the gradients for the selected local coordinate axes are stacked on an additional axis at index 0.

Parameters:

  • input_shape (Tuple[int, ...]) – Shape of input array.

  • axes (Optional[Tuple[int, ...]]) – Axes over which to compute the gradient. Should be a tuple \((i_x, i_y, i_z)\), where \(i_x\), \(i_y\) and \(i_z\) are input array axes assigned to \(x\), \(y\), and \(z\) coordinates respectively. Defaults to None, in which case the axes are taken to be (0, 1, 2). If an integer, this operator returns a jax.Array. If a tuple or None, the resulting arrays are stacked into a BlockArray.

  • center (Union[Tuple[int, ...], Array, None]) – Center of the spherical coordinate system in array indexing coordinates. Default is None, which places the center at the center of the input array.

  • azimuthal (bool) – Flag indicating whether to compute gradients in the azimuthal direction.

  • polar (bool) – Flag indicating whether to compute gradients in the polar direction.

  • radial (bool) – Flag indicating whether to compute gradients in the radial direction.

  • cdiff (bool) – If True, estimate gradients using the second order central different returned by snp.gradient, otherwise use the first order asymmetric difference returned by snp.diff.

  • input_dtype (DType) – dtype for input argument. Default is float32.

  • jit (bool) – If True, jit the evaluation, adjoint, and gram functions of the LinearOperator.

class scico.linop.ComposedLinearOperator(A, B, jit=False)

Bases: LinearOperator

A composition of two LinearOperator objects.

Inheritance diagram of ComposedLinearOperator

A new LinearOperator formed by the composition of two other LinearOperator objects.

A ComposedLinearOperator AB implements AB @ x == A @ B @ x. LinearOperator A and B are stored as attributes of the ComposedLinearOperator.

LinearOperator A and B must have compatible shapes and dtypes: A.input_shape == B.output_shape and A.input_dtype == B.input_dtype.

Parameters:
class scico.linop.LinearOperator(input_shape, output_shape=None, eval_fn=None, adj_fn=None, input_dtype=<class 'numpy.float32'>, output_dtype=None, jit=False)

Bases: Operator

Generic linear operator base class

Inheritance diagram of LinearOperator

Parameters:

  • input_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]) – Shape of input array.

  • output_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...], None]) – Shape of output array. Defaults to None. If None, output_shape is determined by evaluating self.__call__ on an input array of zeros.

  • eval_fn (Optional[Callable]) – Function used in evaluating this LinearOperator. Defaults to None. If None, then self.__call__ must be defined in any derived classes.

  • adj_fn (Optional[Callable]) – Function used to evaluate the adjoint of this LinearOperator. Defaults to None. If None, the adjoint is not set, and the _set_adjoint will be called silently at the first adj call or can be called manually.

  • input_dtype (DType) – dtype for input argument. Defaults to float32. If the LinearOperator implements complex-valued operations, this must be a complex dtype (typically complex64) for correct adjoint and gradient calculation.

  • output_dtype (Optional[DType]) – dtype for output argument. Defaults to None. If None, output_dtype is determined by evaluating self.__call__ on an input array of zeros.

  • jit (bool) – If True, call jit on this LinearOperator to jit the forward, adjoint, and gram functions. Same as calling jit after the LinearOperator is created.

property H: LinearOperator

Hermitian transpose of this LinearOperator.

Return a new LinearOperator that is the Hermitian transpose of this LinearOperator. For a real-valued LinearOperator A (A.input_dtype is float32 or float64), the LinearOperator A.H is equivalent to A.T. For a complex-valued LinearOperator A (A.input_dtype is complex64 or complex128), the LinearOperator A.H implements the adjoint of A : A.H @ y == A.adj(y) == A.conj().T @ y).

For the non-conjugate transpose, see T.

property T: LinearOperator

Transpose of this LinearOperator.

Return a new LinearOperator that implements the transpose of this LinearOperator. For a real-valued LinearOperator A (A.input_dtype is float32 or float64), the LinearOperator A.T implements the adjoint: A.T(y) == A.adj(y). For a complex-valued LinearOperator A (A.input_dtype is complex64 or complex128), the LinearOperator A.T is not the adjoint. For the conjugate transpose, use .conj().T or H.

__call__(x)[source]

Evaluate this LinearOperator at the point \(\mb{x}\).

Parameters:

x (Union[LinearOperator, Array, BlockArray]) – Point at which to evaluate this LinearOperator. If x is a jax.Array or BlockArray, must have shape == self.input_shape. If x is a LinearOperator, must have x.output_shape == self.input_shape.

Return type:

Union[LinearOperator, Array, BlockArray]

adj(y)[source]

Adjoint of this LinearOperator.

Compute the adjoint of this LinearOperator applied to input y.

Parameters:

y (Union[LinearOperator, Array, BlockArray]) – Point at which to compute adjoint. If y is jax.Array or BlockArray, must have shape == self.output_shape. If y is a LinearOperator, must have y.output_shape == self.output_shape.

Return type:

Union[LinearOperator, Array, BlockArray]

Returns:

Adjoint evaluated at y.

conj()[source]

Complex conjugate of this LinearOperator.

Return a new LinearOperator Ac such that Ac(x) = conj(A)(x).

Return type:

LinearOperator

gram(x)[source]

Compute A.adj(A(x)).

Parameters:

x (Union[LinearOperator, Array, BlockArray]) – Point at which to evaluate the gram operator. If x is a jax.Array or BlockArray, must have shape == self.input_shape. If x is a LinearOperator, must have x.output_shape == self.input_shape.

Return type:

Union[LinearOperator, Array, BlockArray]

Returns:

Result of A.adj(A(x)).

property gram_op: LinearOperator

Gram operator of this LinearOperator.

Return a new LinearOperator G such that G(x) = A.adj(A(x))).

jit()[source]

Replace the private functions _eval, _adj, _gram with jitted versions.

class scico.linop.MatrixOperator(A, input_cols=0)

Bases: LinearOperator

Linear operator implementing matrix multiplication.

Inheritance diagram of MatrixOperator

Parameters:

  • A (Union[Array, ndarray, bool, number, bool, int, float, complex]) – Dense array. The action of the created LinearOperator will implement matrix multiplication with A.

  • input_cols (int) – If this parameter is set to the default of 0, the MatrixOperator takes a vector (one-dimensional array) input. If the input is intended to be a matrix (two-dimensional array), this parameter should specify number of columns in the matrix.

property H

Hermitian (conjugate) transpose of this MatrixOperator.

Return a MatrixOperator corresponding to the Hermitian (conjugate) transpose of this matrix.

property T

Transpose of this MatrixOperator.

Return a MatrixOperator corresponding to the transpose of this matrix.

__call__(other)[source]

Evaluate this LinearOperator at the point \(\mb{x}\).

Parameters:

x – Point at which to evaluate this LinearOperator. If x is a jax.Array or BlockArray, must have shape == self.input_shape. If x is a LinearOperator, must have x.output_shape == self.input_shape.

adj(y)[source]

Adjoint of this LinearOperator.

Compute the adjoint of this LinearOperator applied to input y.

Parameters:

y – Point at which to compute adjoint. If y is jax.Array or BlockArray, must have shape == self.output_shape. If y is a LinearOperator, must have y.output_shape == self.output_shape.

Returns:

Adjoint evaluated at y.

conj()[source]

Complex conjugate of this MatrixOperator.

Return a MatrixOperator with complex conjugated elements.

gram(other)[source]

Compute A.adj(A(x)).

Parameters:

x – Point at which to evaluate the gram operator. If x is a jax.Array or BlockArray, must have shape == self.input_shape. If x is a LinearOperator, must have x.output_shape == self.input_shape.

Returns:

Result of A.adj(A(x)).

property gram_op

Gram operator of this MatrixOperator.

Return a new LinearOperator G such that G(x) = A.adj(A(x))).

norm(ord=None, axis=None, keepdims=False)[source]

Compute the norm of the dense matrix self.A.

Call scico.numpy.linalg.norm on the dense matrix self.A.

to_array()[source]

Return a numpy.ndarray containing self.A.

class scico.linop.DiagonalReplicated(op, replicates, input_axis=0, output_axis=None, map_type='auto', **kwargs)

Bases: DiagonalReplicated, LinearOperator

A diagonal stack constructed from a single linear operator.

Inheritance diagram of DiagonalReplicated

Given linear operator \(A\), create the linear operator

\[\begin{split}H = \begin{pmatrix} A & 0 & \ldots & 0\\ 0 & A & \ldots & 0\\ \vdots & \vdots & \ddots & \vdots\\ 0 & 0 & \ldots & A \\ \end{pmatrix} \qquad \text{such that} \qquad H \begin{pmatrix} \mb{x}_1 \\ \mb{x}_2 \\ \vdots \\ \mb{x}_N \\ \end{pmatrix} = \begin{pmatrix} A(\mb{x}_1) \\ A(\mb{x}_2) \\ \vdots \\ A(\mb{x}_N) \\ \end{pmatrix} \;.\end{split}\]

The application of \(A\) to each component \(\mb{x}_k\) is computed using jax.pmap or jax.vmap. The input shape for linear operator \(A\) should exclude the array axis on which \(A\) is replicated to form \(H\). For example, if \(A\) has input shape (3, 4) and \(H\) is constructed to replicate on axis 0 with 2 replicates, the input shape of \(H\) will be (2, 3, 4).

Linear operators taking BlockArray input are not supported.

Parameters:

  • op (LinearOperator) – Linear operator to replicate.

  • replicates (int) – Number of replicates of op.

  • input_axis (int) – Input axis over which op should be replicated.

  • output_axis (Optional[int]) – Index of replication axis in output array. If None, the input replication axis is used.

  • map_type (str) – If “pmap” or “vmap”, apply replicated mapping using jax.pmap or jax.vmap respectively. If “auto”, use jax.pmap if sufficient devices are available for the number of replicates, otherwise use jax.vmap.

class scico.linop.DiagonalStack(ops, collapse_input=True, collapse_output=True, jit=True, **kwargs)

Bases: DiagonalStack, LinearOperator

A diagonal stack of linear operators.

Inheritance diagram of DiagonalStack

Given linear operators \(A_1, A_2, \dots, A_N\), create the linear operator

\[\begin{split}H = \begin{pmatrix} A_1 & 0 & \ldots & 0\\ 0 & A_2 & \ldots & 0\\ \vdots & \vdots & \ddots & \vdots\\ 0 & 0 & \ldots & A_N \\ \end{pmatrix} \qquad \text{such that} \qquad H \begin{pmatrix} \mb{x}_1 \\ \mb{x}_2 \\ \vdots \\ \mb{x}_N \\ \end{pmatrix} = \begin{pmatrix} A_1(\mb{x}_1) \\ A_2(\mb{x}_2) \\ \vdots \\ A_N(\mb{x}_N) \\ \end{pmatrix} \;.\end{split}\]

By default, if the inputs \(\mb{x}_1, \mb{x}_2, \dots, \mb{x}_N\) all have the same (possibly nested) shape, S, this operator will work on the stack, i.e., have an input shape of (N, *S). If the inputs have distinct shapes, S1, S2, …, SN, this operator will work on the block concatenation, i.e., have an input shape of (S1, S2, …, SN). The same holds for the output shape.

Parameters:

  • ops (Sequence[LinearOperator]) – Operators to stack.

  • collapse_input (Optional[bool]) – If True, inputs are expected to be stacked along the first dimension when possible.

  • collapse_output (Optional[bool]) – If True, the output will be stacked along the first dimension when possible.

  • jit (bool) – See jit in LinearOperator.

class scico.linop.VerticalStack(ops, collapse_output=True, jit=True, **kwargs)

Bases: VerticalStack, LinearOperator

A vertical stack of linear operators.

Inheritance diagram of VerticalStack

Given linear operators \(A_1, A_2, \dots, A_N\), create the linear operator

\[\begin{split}H = \begin{pmatrix} A_1 \\ A_2 \\ \vdots \\ A_N \\ \end{pmatrix} \qquad \text{such that} \qquad H \mb{x} = \begin{pmatrix} A_1(\mb{x}) \\ A_2(\mb{x}) \\ \vdots \\ A_N(\mb{x}) \\ \end{pmatrix} \;.\end{split}\]
Parameters:

  • ops (Sequence[LinearOperator]) – Linear operators to stack.

  • collapse_output (Optional[bool]) – If True and the output would be a BlockArray with shape ((m, n, …), (m, n, …), …), the output is instead a jax.Array with shape (S, m, n, …) where S is the length of ops.

  • jit (bool) – See jit in LinearOperator.

scico.linop.linop_over_axes(linop, input_shape, *args, axes=None, **kwargs)

Construct a list of LinearOperator by iterating over axes.

Construct a list of LinearOperator by iterating over a specified sequence of axes, passing each value in sequence to the axis keyword argument of the LinearOperator initializer.

Parameters:
Return type:

List[LinearOperator]

Returns:

A tuple (axes, ops) where axes is a tuple of the axes used to construct the list of LinearOperator, and ops is the list itself.

scico.linop.jacobian(F, u, include_eval=False)

Construct Jacobian linear operator for a general operator.

For a specified Operator, construct a corresponding Jacobian LinearOperator, the application of which is equivalent to multiplication by the Jacobian of the Operator at a specified input value.

The implementation of this function is based on Operator.jvp and Operator.vjp, which are themselves based on jax.jvp and jax.vjp. For reasons of computational efficiency, these functions return the value of the Operator evaluated at the specified point in addition to the requested Jacobian-vector product. If the include_eval parameter of this function is True, the constructed LinearOperator returns a BlockArray output, the first component of which is the result of the Operator evaluation, and the second component of which is the requested Jacobian-vector product. If include_eval is False, then the Operator evaluation computed by jax.jvp and jax.vjp are discarded.

Parameters:

  • F (Operator) – Operator of which the Jacobian is to be computed.

  • u (Array) – Input value of the Operator at which the Jacobian is to be computed.

  • include_eval (Optional[bool]) – Flag indicating whether the result of evaluating the Operator should be included (as the first component of a BlockArray) in the output of the Jacobian LinearOperator constructed by this function.

Return type:

LinearOperator

Returns:

A LinearOperator capable of computing Jacobian-vector products.

scico.linop.operator_norm(A, maxiter=100, key=None)

Estimate the norm of a LinearOperator.

Estimate the operator norm induced by the \(\ell_2\) vector norm, i.e. for LinearOperator \(A\),

\[\begin{split}\| A \|_2 &= \max \{ \| A \mb{x} \|_2 \, : \, \| \mb{x} \|_2 \leq 1 \} \\ &= \sqrt{ \lambda_{ \mathrm{max} }( A^H A ) } = \sigma_{\mathrm{max}}(A) \;,\end{split}\]

where \(\lambda_{\mathrm{max}}(B)\) and \(\sigma_{\mathrm{max}}(B)\) respectively denote the largest eigenvalue of \(B\) and the largest singular value of \(B\). The value is estimated via power iteration, using power_iteration, to estimate \(\lambda_{\mathrm{max}}(A^H A)\).

Parameters:

  • A (LinearOperator) – LinearOperator for which operator norm is desired.

  • maxiter (int) – Maximum number of power iterations to use. Default: 100

  • key (Optional[Array]) – Jax PRNG key. Defaults to None, in which case a new key is created.

Returns:

float – Norm of operator \(A\).

scico.linop.power_iteration(A, maxiter=100, key=None)

Compute largest eigenvalue of a diagonalizable LinearOperator.

Compute largest eigenvalue of a diagonalizable LinearOperator using power iteration.

Parameters:

  • A (LinearOperator) – LinearOperator used for computation. Must be diagonalizable.

  • maxiter (int) – Maximum number of power iterations to use.

  • key (Optional[Array]) – Jax PRNG key. Defaults to None, in which case a new key is created.

Returns:

tuple –

A tuple (mu, v) containing:

  • mu: Estimate of largest eigenvalue of A.

  • v: Eigenvector of A with eigenvalue mu.

scico.linop.valid_adjoint(A, AT, eps=1e-07, x=None, y=None, key=None)

Check whether LinearOperator AT is the adjoint of A.

Check whether LinearOperator \(\mathsf{AT}\) is the adjoint of \(\mathsf{A}\). The test exploits the identity

\[\mathbf{y}^T (A \mathbf{x}) = (\mathbf{y}^T A) \mathbf{x} = (A^T \mathbf{y})^T \mathbf{x}\]

by computing \(\mathbf{u} = \mathsf{A}(\mathbf{x})\) and \(\mathbf{v} = \mathsf{AT}(\mathbf{y})\) for random \(\mathbf{x}\) and \(\mathbf{y}\) and confirming that

\[\frac{| \mathbf{y}^T \mathbf{u} - \mathbf{v}^T \mathbf{x} |} {\max \left\{ | \mathbf{y}^T \mathbf{u} |, | \mathbf{v}^T \mathbf{x} | \right\}} < \epsilon \;.\]

If \(\mathsf{A}\) is a complex operator (with a complex input_dtype) then the test checks whether \(\mathsf{AT}\) is the Hermitian conjugate of \(\mathsf{A}\), with a test as above, but with all the \((\cdot)^T\) replaced with \((\cdot)^H\).

Parameters:

  • A (LinearOperator) – Primary LinearOperator.

  • AT (LinearOperator) – Adjoint LinearOperator.

  • eps (Optional[float]) – Error threshold for validation of \(\mathsf{AT}\) as adjoint of \(\mathsf{AT}\). If None, the relative error is returned instead of a boolean value.

  • x (Optional[Array]) – If not the default None, use the specified array instead of a random array as test vector \(\mb{x}\). If specified, the array must have shape A.input_shape.

  • y (Optional[Array]) – If not the default None, use the specified array instead of a random array as test vector \(\mb{y}\). If specified, the array must have shape AT.input_shape.

  • key (Optional[Array]) – Jax PRNG key. Defaults to None, in which case a new key is created.

Return type:

Union[bool, float]

Returns:

Boolean value indicating whether validation passed, or relative error of test, depending on type of parameter eps.