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 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(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(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 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))