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