Linear Operators

class cola.ops.Dense(A)

Bases: LinearOperator

LinearOperator wrapping of a dense matrix. O(n^2) memory and time mvms.

Parameters:

A (array_like) – Dense matrix to be wrapped.

Example

>>> A = jnp.array([[1., 2.], [3., 4.]])
>>> op = Dense(A)

to_dense()

Produces a dense array representation of the linear operator.

class cola.ops.Triangular(A, lower=True)

Bases: Dense

Triangular Linear Operator.

class cola.ops.Sparse(*args, **kwargs)

Bases: LinearOperator

Sparse linear operator.

Parameters:

• data (array_like) – 1-D array representing the nonzero values of the sparse matrix.

• row_indices (array_like) – 1-D array representing the row indices of the nonzero values.

• col_indices (array_like) – 1-D array representing the column indices of the nonzero values.

• shape (tuple) – Shape of the sparse matrix.

Example

>>> data = jnp.array([1, 2, 3, 4, 5, 6])
>>> rol_indices = jnp.array([0, 0, 1, 2, 2, 2])
>>> col_indices = jnp.array([1, 3, 3, 0, 1, 2])
>>> shape = (3, 4)
>>> op = Sparse(data, row_indices, col_indices, shape)

class cola.ops.ScalarMul(*args, **kwargs)

Bases: LinearOperator

Linear Operator representing scalar multiplication

class cola.ops.Identity(*args, **kwargs)

Bases: LinearOperator

Linear Operator representing the identity matrix. Can also be created from I_like(A)

Parameters:

• shape (tuple) – Shape of the identity matrix.

• dtype – Data type of the identity matrix.

Example

>>> shape = (3, 3)
>>> dtype =  jnp.float64
>>> op = Identity(shape, dtype)

to(device)

Returns a new linear operator with given device and dtype WARNING: dtype change is not supported yet.

class cola.ops.Product(*args, **kw_args)

Bases: Product

Matrix Multiply Product of Linear ops

class cola.ops.Sum(*args, **kw_args)

Bases: Sum

Sum of Linear ops

class cola.ops.Kronecker(*args, **kw_args)

Bases: Kronecker

Kronecker product of linear ops Kronecker([M1,M2]):= M1⊗M2

Parameters:

*Ms (array_like) – Sequence of linear operators representing the Kronecker product operands.

Example

>>> M1 = jnp.array([[1, 2], [3, 4]])
>>> M2 = jnp.array([[5, 6], [7, 8]])
>>> op = Kronecker(M1, M2)

class cola.ops.KronSum(*args, **kw_args)

Bases: KronSum

Kronecker Sum Linear Operator, KronSum(A,B):= A ⊕ B = A ⊗ I + I ⊗ B

Parameters:

*Ms (array_like) – Sequence of matrices representing the Kronecker sum operands.

Example

>>> M1 = jnp.array([[1, 2], [3, 4]])
>>> M2 = jnp.array([[5, 6], [7, 8]])
>>> op = KronSum(M1, M2)

class cola.ops.BlockDiag(*args, **kw_args)

Bases: BlockDiag

Block Diagonal Linear Operator. BlockDiag([A,B]):= [A 0; 0 B]

Parameters:

• *Ms (array_like) – Sequence of matrices representing the blocks.

• multiplicities (list, optional) – List of integers representing the multiplicities of the corresponding blocks in *Ms. Default is None, which assigns a multiplicity of 1 to each block.

Example

>>> M1 = jnp.array([[1, 2], [3, 4]])
>>> M2 = jnp.array([[5, 6], [7, 8]])
>>> op = BlockDiag(M1, M2, multiplicities=[2, 3])

class cola.ops.Diagonal(*args, **kwargs)

Bases: LinearOperator

Diagonal LinearOperator. O(n) time and space matmuls.

Parameters:

diag (array_like) – 1-D array representing the diagonal elements of the matrix.

Example

>>> d = jnp.array([1, 2, 3])
>>> op = Diagonal(d)

to_dense()

Produces a dense array representation of the linear operator.

class cola.ops.Tridiagonal(alpha, beta, gamma)

Bases: LinearOperator

Tridiagonal linear operator. O(n) time and space matmuls.

Parameters:

• alpha (array_like) – 1-D array representing lower band of the operator.

• beta (array_like) – 1-D array representing diagonal of the operator.

• gamma (array_like) – 1-D array representing upper band of the operator.

class cola.ops.Transpose(*args, **kw_args)

Bases: Transpose

Transpose of a Linear Operator

class cola.ops.Adjoint(*args, **kw_args)

Bases: Adjoint

Complex conjugate transpose of a Linear Operator (aka adjoint)

class cola.ops.Sliced(*args, **kw_args)

Bases: Sliced

Slicing of another linear operator A. Equivalent to A[slices[0], :][:, slices[1]]

class cola.ops.Jacobian(*args, **kwargs)

Bases: LinearOperator

Jacobian (linearization) of a function f: R^n -> R^m at point x.

Matrix has shape (m, n)

Parameters:

• f (callable) – Function representing the mapping from R^n to R^m.

• x (array_like) – 1-D array representing the point at which to compute the Jacobian.

Example

>>> def f(x):
...     return  jnp.array([x[0]**2, x[1]**3, jnp.sin(x[2])])
>>> x =  jnp.array([1, 2, 3])
>>> op = Jacobian(f, x)

class cola.ops.Hessian(*args, **kwargs)

Bases: LinearOperator

Hessian of a scalar function f: R^n -> R at point x.

Matrix has shape (n, n)

Parameters:

• f (callable) – Function representing the mapping from R^n to R.

• x (array_like) – 1-D array representing the point at which to compute the Hessian.

Example

>>> def f(x):
...     return x[1]**3+np.sin(x[2])
>>> x =  jnp.array([1, 2, 3])
>>> op = Hessian(f, x)

class cola.ops.Permutation(*args, **kwargs)

Bases: LinearOperator

Permutation matrix.

Parameters:

• perm (array_like) – 1-D array representing the permutation.

• dtype (optional) – specify the dtype to operate on (not int)

Example

>>> P = Permutation(np.array([1, 0, 3, 2]))

class cola.ops.Concatenated(*args, **kw_args)

Bases: Concatenated

Produces a linear operator equivalent to concatenating

a collection of matrices Ms along specified axis

Parameters:

• *Ms (array_like) – Sequence of matrices representing the blocks.

• axis (int, optional) – specify which axis to concatenate on (0 or 1)

Example

>>> M1 = jnp.array([[1, 2], [3, 4]])
>>> M2 = jnp.array([[5, 6], [7, 8]])
>>> A = Concatenated(M1, M2, axis=1)
>>> A.shape
>>> (2,4)

class cola.ops.ConvolveND(*args, **kwargs)

Bases: LinearOperator

n-Dimensional convolution Linear operator (only works in jax right now.)

class cola.ops.Householder(*args, **kwargs)

Bases: LinearOperator

Householder rotation matrix.

class cola.ops.Kernel(*args, **kwargs)

Bases: LinearOperator

Kernel operator based on a given function f where the matvec is evaluated on the fly. That is, [Kv]_i = sum_{j} f(x1_i, x2_j) v_j. The variables block_size1 and block_size2 determine the memory usage of the matvec and matmat operations.

Args:

x1 (array): N-D array x2 (array): N-D array fn (callable): function that defines the kernel block_size1 (int): block size for x1 block_size2 (int): block size for x2

cola.ops.I_like(A)

A function that produces an Identity operator with the same shape, dtype and device as A

Return type:

Identity

class cola.ops.FFT(*args, **kwargs)

Bases: LinearOperator

FFT matrix. Uses convention so matrix is unitary.