def collapse(bmat):
'''Collapse what are blocks of bmat'''
# Single block cases
# Do nothing
if is_petsc_mat(bmat) or is_number(bmat) or is_petsc_vec(bmat):
return bmat
# Multiplication
if isinstance(bmat, block_mul):
return collapse_mul(bmat)
# +
elif isinstance(bmat, block_add):
return collapse_add(bmat)
# -
elif isinstance(bmat, block_sub):
return collapse_sub(bmat)
# T
elif isinstance(bmat, block_transpose):
return collapse_tr(bmat)
# Some things in cbc.block know their matrix representation
# E.g. InvLumpDiag...
elif hasattr(bmat, 'A'):
assert is_petsc_mat(bmat.A)
return bmat.A
raise ValueError('Do not know how to collapse %r' % type(bmat))
def collapse_mul(bmat):
'''A*B*C to single matrix'''
# A0 * A1 * ...
A, B = bmat.chain[0], bmat.chain[1:]
if len(B) == 1:
B = B[0]
# Two matrices
if is_petsc_mat(A) and is_petsc_mat(B):
A_ = as_petsc(A)
B_ = as_petsc(B)
assert A_.size[1] == B_.size[0]
C_ = PETSc.Mat()
A_.matMult(B_, C_)
return PETScMatrix(C_)
# One of them is a number
elif is_petsc_mat(A) and is_number(B):
A_ = as_petsc(A)
C_ = A_.copy()
C_.scale(B)
return PETScMatrix(C_)
elif is_petsc_mat(B) and is_number(A):
B_ = as_petsc(B)
C_ = B_.copy()
C_.scale(A)
return PETScMatrix(C_)
# Some compositions
else:
return collapse(collapse(A)*collapse(B))
# Recurse
else:
return collapse_mul(collapse(A)*collapse(reduce(operator.mul, B)))
def collapse_mul(bmat):
'''A*B*C to single matrix'''
# A0 * A1 * ...
A, B = bmat.chain[0], bmat.chain[1:]
if len(B) == 1:
B = B[0]
# Two matrices
if is_petsc_mat(A) and is_petsc_mat(B):
A_ = as_petsc(A)
B_ = as_petsc(B)
assert A_.size[1] == B_.size[0]
C_ = PETSc.Mat()
A_.matMult(B_, C_)
return PETScMatrix(C_)
# One of them is a number
elif is_petsc_mat(A) and is_number(B):
A_ = as_petsc(A)
C_ = A_.copy()
C_.scale(B)
return PETScMatrix(C_)
elif is_petsc_mat(B) and is_number(A):
B_ = as_petsc(B)
C_ = B_.copy()
C_.scale(A)
return PETScMatrix(C_)
# Some compositions
else:
return collapse(collapse(A) * collapse(B))
# Recurse
else:
return collapse_mul(collapse(A) * collapse(reduce(operator.mul, B)))
def get_dims(thing):
'''
Size of Rn vector or operator Rn to Rm. We return None for scalars
and raise when such an operator cannot be established, i.e. there
are consistency checks going on
'''
if is_petsc_vec(thing): return thing.size()
if is_petsc_mat(thing): return (thing.size(0), thing.size(1))
if is_number(thing): return None
# Now let's handdle block stuff
# Multiplication
if isinstance(thing, block_mul):
A, B = thing.chain[0], thing.chain[1:]
dims_A, dims_B = get_dims(A), get_dims(B[0])
# A number does not change
if dims_A is None:
return dims_B
if dims_B is None:
return dims_A
# Otherwise, consistency
if len(B) == 1:
assert len(dims_A) == len(dims_B)
assert dims_A[1] == dims_B[0]
return (dims_A[0], dims_B[1])
else:
dims_B = get_dims(reduce(operator.mul, B))
assert len(dims_A) == len(dims_B)
assert dims_A[1] == dims_B[0]
return (dims_A[0], dims_B[1])
# +, -
elif isinstance(thing, (block_add, block_sub)):
A, B = thing.A, thing.B
if is_number(A):
return get_dims(B)
if is_number(B):
return get_dims(A)
dims = get_dims(A)
assert dims == get_dims(B), (dims, get_dims(B))
return dims
# T
elif isinstance(thing, block_transpose):
dims = get_dims(thing.A)
return (dims[1], dims[0])
# Some things in cbc.block know their matrix representation
# E.g. InvLumpDiag...
elif hasattr(thing, 'A'):
assert is_petsc_mat(thing.A)
return get_dims(thing.A)
raise ValueError('Cannot get_dims of %r, %s' % (type(thing), thing))
def collapse_sub(bmat):
'''A - B to single matrix'''
A, B = bmat.A, bmat.B
# Base case
if is_petsc_mat(A) and is_petsc_mat(B):
A_ = as_petsc(A)
B_ = as_petsc(B)
assert A_.size == B_.size
C_ = A_.copy()
# C = A - B
C_.axpy(-1., B_, PETSc.Mat.Structure.DIFFERENT)
return PETScMatrix(C_)
# Recurse
return collapse_sub(collapse(A) - collapse(B))
def collapse_sub(bmat):
'''A - B to single matrix'''
A, B = bmat.A, bmat.B
# Base case
if is_petsc_mat(A) and is_petsc_mat(B):
A_ = as_petsc(A)
B_ = as_petsc(B)
assert A_.size == B_.size
C_ = A_.copy()
# C = A - B
C_.axpy(-1., B_, PETSc.Mat.Structure.DIFFERENT)
return PETScMatrix(C_)
# Recurse
return collapse_sub(collapse(A) - collapse(B))
def convert(bmat, algorithm='numpy'):
'''
Attempt to convert bmat to a PETSc(Matrix/Vector) object.
If succed this is at worst a number.
'''
# Block vec conversion
if isinstance(bmat, block_vec):
array = block_vec_to_numpy(bmat)
vec = PETSc.Vec().createWithArray(array)
vec.assemble()
return PETScVector(vec)
# Conversion of bmat is bit more involved because of the possibility
# that some of the blocks are numbers or composition of matrix operations
if isinstance(bmat, block_mat):
# Create collpsed bmat
row_sizes, col_sizes = bmat_sizes(bmat)
nrows, ncols = len(row_sizes), len(col_sizes)
indices = itertools.product(list(range(nrows)), list(range(ncols)))
blocks = np.zeros((nrows, ncols), dtype='object')
for block, (i, j) in zip(bmat.blocks.flatten(), indices):
# This might is guaranteed to be matrix or number
A = collapse(block)
if is_number(A):
# Diagonal matrices can be anything provided square
if i == j and row_sizes[i] == col_sizes[j]:
A = diagonal_matrix(row_sizes[i], A)
else:
# Offdiagonal can only be zero
A = zero_matrix(row_sizes[i], col_sizes[j])
#else:
# A = 0
# The converted block
blocks[i, j] = A
# Now every block is a matrix/number and we can make a monolithic thing
bmat = block_mat(blocks)
assert all(is_petsc_mat(block) or is_number(block)
for block in bmat.blocks.flatten())
# Opt out of monolithic
if not algorithm:
set_lg_map(bmat)
return bmat
# Monolithic via numpy (fast)
# Convert to numpy
array = block_mat_to_numpy(bmat)
# Constuct from numpy
bmat = numpy_to_petsc(array)
set_lg_map(bmat)
return bmat
# Try with a composite
return collapse(bmat)
def collapse(bmat):
'''Collapse what are blocks of bmat'''
# Single block cases
# Do nothing
if is_petsc_mat(bmat) or is_number(bmat) or is_petsc_vec(bmat):
return bmat
if isinstance(bmat, (Vector, Matrix, GenericVector)):
return bmat
# Multiplication
if isinstance(bmat, block_mul):
return collapse_mul(bmat)
# +
elif isinstance(bmat, block_add):
return collapse_add(bmat)
# -
elif isinstance(bmat, block_sub):
return collapse_sub(bmat)
# T
elif isinstance(bmat, block_transpose):
return collapse_tr(bmat)
# Some things in cbc.block know their matrix representation
# This is typically diagonals like InvLumpDiag etc
elif hasattr(bmat, 'v'):
# So now we make that diagonal matrix
diagonal = bmat.v
n = diagonal.size
mat = PETSc.Mat().createAIJ(comm=COMM, size=[[n, n], [n, n]], nnz=1)
mat.assemblyBegin()
mat.setDiagonal(diagonal)
mat.assemblyEnd()
return PETScMatrix(mat)
# Some operators actually have matrix repre (HsMG)
elif hasattr(bmat, 'matrix'):
return bmat.matrix
# Try:
elif hasattr(bmat, 'collapse'):
return bmat.collapse()
elif hasattr(bmat, 'create_vec'):
x = bmat.create_vec()
columns = []
for ei in Rn_basis(x):
y = bmat*ei
columns.append(csr_matrix(convert(y).get_local()))
bmat = (sp_vstack(columns).T).tocsr()
return numpy_to_petsc(bmat)
raise ValueError('Do not know how to collapse %r' % type(bmat))
def collapse_tr(bmat):
'''to Transpose'''
# Base
A = bmat.A
if is_petsc_mat(A):
A_ = as_petsc(A)
C_ = PETSc.Mat()
A_.transpose(C_)
return PETScMatrix(C_)
# Recurse
return collapse_tr(collapse(bmat))
def collapse_tr(bmat):
'''to Transpose'''
# Base
A = bmat.A
if is_petsc_mat(A):
A_ = as_petsc(A)
C_ = PETSc.Mat()
A_.transpose(C_)
return PETScMatrix(C_)
# Recurse
return collapse_tr(collapse(bmat))
def convert(bmat, algorithm='numpy'):
'''
Attempt to convert bmat to a PETSc(Matrix/Vector) object.
If succed this is at worst a number.
'''
# Block vec conversion
if isinstance(bmat, block_vec):
array = block_vec_to_numpy(bmat)
vec = PETSc.Vec().createWithArray(array)
vec.assemble()
return PETScVector(vec)
# Conversion of bmat is bit more involved because of the possibility
# that some of the blocks are numbers or composition of matrix operations
if isinstance(bmat, block_mat):
# Create collpsed bmat
row_sizes, col_sizes = bmat_sizes(bmat)
nrows, ncols = len(row_sizes), len(col_sizes)
indices = itertools.product(range(nrows), range(ncols))
blocks = np.zeros((nrows, ncols), dtype='object')
for block, (i, j) in zip(bmat.blocks.flatten(), indices):
# This might is guaranteed to be matrix or number
A = collapse(block)
if is_number(A):
# Diagonal matrices can be anything provided square
if i == j and row_sizes[i] == col_sizes[j]:
A = diagonal_matrix(row_sizes[i], A)
else:
# Offdiagonal can only be zero
A = zero_matrix(row_sizes[i], col_sizes[j])
#else:
# A = 0
# The converted block
blocks[i, j] = A
# Now every block is a matrix/number and we can make a monolithic thing
bmat = block_mat(blocks)
assert all(is_petsc_mat(block) or is_number(block)
for block in bmat.blocks.flatten())
# Opt out of monolithic
if not algorithm: return bmat
# Monolithic via numpy (fast)
# Convert to numpy
array = block_mat_to_numpy(bmat)
# Constuct from numpy
return numpy_to_petsc(array)
# Try with a composite
return collapse(bmat)
def block_mat_to_numpy(bmat):
'''Collapsing block mat of matrices to scipy's bmat'''
# A single matrix
if is_petsc_mat(bmat):
bmat = as_petsc(bmat)
return csr_matrix(bmat.getValuesCSR()[::-1], shape=bmat.size)
# 0
if is_number(bmat):
return None # What bmat accepts
# Recurse on blocks
blocks = np.array(map(block_mat_to_numpy, bmat.blocks.flatten()))
blocks = blocks.reshape(bmat.blocks.shape)
# The bmat
return numpy_block_mat(blocks).tocsr()
def block_mat_to_numpy(bmat):
'''Collapsing block mat of matrices to scipy's bmat'''
# A single matrix
if is_petsc_mat(bmat):
bmat = as_petsc(bmat)
return csr_matrix(bmat.getValuesCSR()[::-1], shape=bmat.size)
# 0
if is_number(bmat):
return None # What bmat accepts
# Recurse on blocks
blocks = np.array(map(block_mat_to_numpy, bmat.blocks.flatten()))
blocks = blocks.reshape(bmat.blocks.shape)
# The bmat
return numpy_block_mat(blocks).tocsr()
def collapse(bmat):
'''Collapse what are blocks of bmat'''
# Single block cases
# Do nothing
if is_petsc_mat(bmat) or is_number(bmat) or is_petsc_vec(bmat):
return bmat
if isinstance(bmat, (Vector, Matrix, GenericVector)):
return bmat
# Multiplication
if isinstance(bmat, block_mul):
return collapse_mul(bmat)
# +
elif isinstance(bmat, block_add):
return collapse_add(bmat)
# -
elif isinstance(bmat, block_sub):
return collapse_sub(bmat)
# T
elif isinstance(bmat, block_transpose):
return collapse_tr(bmat)
# Some things in cbc.block know their matrix representation
# This is typically diagonals like InvLumpDiag etc
elif hasattr(bmat, 'v'):
# So now we make that diagonal matrix
diagonal = bmat.v
n = diagonal.size
mat = PETSc.Mat().createAIJ(comm=COMM, size=[[n, n], [n, n]], nnz=1)
mat.assemblyBegin()
mat.setDiagonal(diagonal)
mat.assemblyEnd()
return PETScMatrix(mat)
# Some operators actually have matrix repre (HsMG)
elif hasattr(bmat, 'matrix'):
return bmat.matrix
raise ValueError('Do not know how to collapse %r' % type(bmat))
def collapse(bmat):
'''Collapse what are blocks of bmat'''
# Single block cases
# Do nothing
if is_petsc_mat(bmat) or is_number(bmat) or is_petsc_vec(bmat):
return bmat
if isinstance(bmat, (Vector, Matrix, GenericVector)):
return bmat
# Multiplication
if isinstance(bmat, block_mul):
return collapse_mul(bmat)
# +
elif isinstance(bmat, block_add):
return collapse_add(bmat)
# -
elif isinstance(bmat, block_sub):
return collapse_sub(bmat)
# T
elif isinstance(bmat, block_transpose):
return collapse_tr(bmat)
# Some things in cbc.block know their matrix representation
# This is typically diagonals like InvLumpDiag etc
elif hasattr(bmat, 'v'):
# So now we make that diagonal matrix
diagonal = bmat.v
n = diagonal.size
mat = PETSc.Mat().createAIJ(comm=COMM, size=[[n, n], [n, n]], nnz=1)
mat.assemblyBegin()
mat.setDiagonal(diagonal)
mat.assemblyEnd()
return PETScMatrix(mat)
# Some operators actually have matrix repre (HsMG)
elif hasattr(bmat, 'matrix'):
return bmat.matrix
raise ValueError('Do not know how to collapse %r' % type(bmat))
def set_lg_map(mat):
'''Set local-to-global-map on the matrix'''
# NOTE; serial only - so we own everything but still sometimes we need
# to tell that to petsc (especiialy when bcs are to be applied)
if is_number(mat): return mat
assert is_petsc_mat(mat) or isinstance(mat, block_mat), (type(mat))
if isinstance(mat, block_mat):
blocks = np.array(map(set_lg_map, mat.blocks.flatten())).reshape(mat.blocks.shape)
return block_mat(blocks)
comm = mpi_comm_world().tompi4py()
# Work with matrix
rowmap, colmap = range(mat.size(0)), range(mat.size(1))
row_lgmap = PETSc.LGMap().create(rowmap, comm=comm)
col_lgmap = PETSc.LGMap().create(colmap, comm=comm)
as_petsc(mat).setLGMap(row_lgmap, col_lgmap)
return mat
def set_lg_map(mat):
'''Set local-to-global-map on the matrix'''
# NOTE; serial only - so we own everything but still sometimes we need
# to tell that to petsc (especiialy when bcs are to be applied)
if is_number(mat): return mat
assert is_petsc_mat(mat) or isinstance(mat, block_mat), (type(mat))
if isinstance(mat, block_mat):
blocks = np.array(map(set_lg_map, mat.blocks.flatten())).reshape(mat.blocks.shape)
return block_mat(blocks)
comm = mpi_comm_world().tompi4py()
# Work with matrix
rowmap, colmap = range(mat.size(0)), range(mat.size(1))
row_lgmap = PETSc.LGMap().create(rowmap, comm=comm)
col_lgmap = PETSc.LGMap().create(colmap, comm=comm)
as_petsc(mat).setLGMap(row_lgmap, col_lgmap)
return mat
def get_dims(thing):
'''
Size of Rn vector or operator Rn to Rm. We return None for scalars
and raise when such an operator cannot be established, i.e. there
are consistency checks going on
'''
if is_petsc_vec(thing): return thing.size()
if is_petsc_mat(thing): return (thing.size(0), thing.size(1))
if is_number(thing): return None
# Now let's handdle block stuff
# Multiplication
if isinstance(thing, block_mul):
A, B = thing.chain[0], thing.chain[1:]
dims_A, dims_B = get_dims(A), get_dims(B[0])
# A number does not change
if dims_A is None:
return dims_B
if dims_B is None:
return dims_A
# Otherwise, consistency
if len(B) == 1:
assert len(dims_A) == len(dims_B)
assert dims_A[1] == dims_B[0], (dims_A, dims_B)
return (dims_A[0], dims_B[1])
else:
dims_B = get_dims(reduce(operator.mul, B))
assert len(dims_A) == len(dims_B)
assert dims_A[1] == dims_B[0], (dims_A, dims_B)
return (dims_A[0], dims_B[1])
# +, -
if isinstance(thing, (block_add, block_sub)):
A, B = thing.A, thing.B
if is_number(A):
return get_dims(B)
if is_number(B):
return get_dims(A)
dims = get_dims(A)
assert dims == get_dims(B), (dims, get_dims(B))
return dims
# T
if isinstance(thing, block_transpose):
dims = get_dims(thing.A)
return (dims[1], dims[0])
# Some things in cbc.block know their maE.g. InvLumpDiag...Almost last resort
if hasattr(thing, 'A'):
assert is_petsc_mat(thing.A)
return get_dims(thing.A)
if hasattr(thing, '__sizes__'):
return thing.__sizes__
if hasattr(thing, 'create_vec'):
return (thing.create_vec(0).size(), thing.create_vec(1).size())
raise ValueError('Cannot get_dims of %r, %s' % (type(thing), thing))
