def plot_matrix_hybrid(A, **kwargs): """Provides a plot of a hybrid-mixed matrix.""" fig, ax = plt.subplots(1, 1) petsc_mat = A.M.handle total_size = petsc_mat.getSize() f0_size = A.M[0, 0].handle.getSize() f1_size = A.M[1, 1].handle.getSize() Mij = PETSc.Mat() petsc_mat.convert("aij", Mij) n, m = total_size Mnp = np.array(Mij.getValues(range(n), range(m))) Am = np.ma.masked_values(Mnp, 0, rtol=1e-13) my_cmap = plt.cm.magma my_cmap.set_bad(color="lightgray") # Plot the matrix plot = ax.matshow(Am, cmap=my_cmap, **kwargs) # plot = plt.spy(Am, **kwargs) # Remove axis ticks and values ax.tick_params(length=0) ax.set_xticklabels([]) ax.set_yticklabels([]) ax.axhline(y=f0_size[0] - 0.5, color="k") ax.axvline(x=f0_size[0] - 0.5, color="k") ax.axhline(y=f0_size[0] + f1_size[0] - 0.5, color="k") ax.axvline(x=f0_size[0] + f1_size[0] - 0.5, color="k") return plot
def create_injection(self, dmc, dmf): prefix = dmc.getOptionsPrefix() mat_type = PETSc.Options(prefix).getString( "mg_levels_transfer_mat_type", default="matfree") I = self.create_transfer(get_appctx(dmf), get_appctx(dmc), mat_type, False, False) return PETSc.Mat().createTranspose(I)
def construct_kronecker_matrix(self, interp_1d): """ Construct the tensorized interpolation matrix. Do this by computing the kron product of the rows of the 1d univariate interpolation matrices. In the future, this may be done matrix-free. """ # this is awfully slow in 3D!!!!!!!!!!! (FW: it might not be anymore) IFW = PETSc.Mat().create(self.comm) IFW.setType(PETSc.Mat.Type.AIJ) comm = self.comm # owned part of global problem local_N = self.N // comm.size + int(comm.rank < (self.N % comm.size)) (lsize, gsize) = interp_1d[0].getSizes()[0] IFW.setSizes(((lsize, gsize), (local_N, self.N))) IFW.setUp() for row in range(lsize): row = self.lg_map_fe.apply([row])[0] denserows = [A[row, :] for A in interp_1d] values = reduce(np.kron, denserows) columns = np.where(values != 0)[0].astype(np.int32) values = values[columns] IFW.setValues([row], columns, values) IFW.assemble() return IFW
def plot_matrix(A, **kwargs): """Provides a plot of a matrix.""" fig, ax = plt.subplots(1, 1) petsc_mat = A.M.handle size = petsc_mat.getSize() Mij = PETSc.Mat() petsc_mat.convert("aij", Mij) n, m = size Mnp = np.array(Mij.getValues(range(n), range(m))) Am = np.ma.masked_values(Mnp, 0, rtol=1e-13) my_cmap = plt.cm.magma my_cmap.set_bad(color="lightgray") # Plot the matrix plot = ax.matshow(Am, cmap=my_cmap, **kwargs) # Remove axis ticks and values ax.tick_params(length=0) ax.set_xticklabels([]) ax.set_yticklabels([]) return plot
def assemble(self): """ Compute values of G and assemble the sparse matrix This method creates the underlying sparse covariance matrix based on the pre-computed entries, allocates memory, sets the values, and finalizes assembly. No inputs, no return value. If this method is called and the matrix has already been assembled, it will return with no effect. """ if self.is_assembled: return G_dict, nnz = self._compute_G_vals() self.G = PETSc.Mat().create(comm=self.comm) self.G.setType('aij') self.G.setSizes(((self.nx_local, -1), (self.nx_local, -1))) self.G.setPreallocationNNZ(nnz) self.G.setFromOptions() self.G.setUp() for key, val in G_dict.items(): self.G.setValues(np.array(key, dtype=PETSc.IntType), val[1], val[0]) self.G.assemble() self.is_assembled = True
def construct_full_interpolation_matrix(self, IFW): """ Assemble interpolation matrix for vectorial tensorized spline space. """ # this is THE MOST awfully slow in 3D!!!!!!!!!!! FullIFW = PETSc.Mat().create(self.comm) FullIFW.setType(PETSc.Mat.Type.AIJ) d = self.dim free_dims = list(set(range(self.dim)) - set(self.fixed_dims)) dfree = len(free_dims) ((lsize, gsize), (lsize_spline, gsize_spline)) = IFW.getSizes() FullIFW.setSizes(((d * lsize, d * gsize), (dfree * lsize_spline, dfree * gsize_spline))) # BIG TODO: figure out the sparsity pattern FullIFW.setUp() # this blows up the matrix to do the right thing # on vector fields. It's not just a block matrix, # but the values are interleaved as this is how # firedrake handles vector fields for row in range(lsize): row = self.lg_map_fe.apply([row])[0] (cols, vals) = IFW.getRow(row) for j, dim in enumerate(free_dims): FullIFW.setValues([d * row + dim], [dfree * col + j for col in cols], vals) FullIFW.assemble() return FullIFW
def createSubMatrix(self, mat, row_is, col_is, target=None): if target is not None: # Repeat call, just return the matrix, since we don't # actually assemble in here. target.assemble() return target # These are the sets of ISes of which the the row and column # space consist. row_ises = self._y.function_space().dof_dset.field_ises col_ises = self._x.function_space().dof_dset.field_ises row_inds = find_sub_block(row_is, row_ises) if row_is == col_is and row_ises == col_ises: col_inds = row_inds else: col_inds = find_sub_block(col_is, col_ises) splitter = ExtractSubBlock() asub = splitter.split(self.a, argument_indices=(row_inds, col_inds)) Wrow = asub.arguments()[0].function_space() Wcol = asub.arguments()[1].function_space() row_bcs = [] col_bcs = [] for bc in self.bcs: if isinstance(bc, DirichletBC): bc_temp = bc.reconstruct(field=row_inds, V=Wrow, g=bc.function_arg, sub_domain=bc.sub_domain, method=bc.method, use_split=True) elif isinstance(bc, EquationBCSplit): bc_temp = bc.reconstruct(field=row_inds, V=Wrow, row_field=row_inds, col_field=col_inds, use_split=True) if bc_temp is not None: row_bcs.append(bc_temp) if Wrow == Wcol and row_inds == col_inds and self.bcs == self.bcs_col: col_bcs = row_bcs else: for bc in self.bcs_col: if isinstance(bc, DirichletBC): bc_temp = bc.reconstruct(field=col_inds, V=Wcol, g=bc.function_arg, sub_domain=bc.sub_domain, method=bc.method, use_split=True) elif isinstance(bc, EquationBCSplit): bc_temp = bc.reconstruct(field=col_inds, V=Wcol, row_field=row_inds, col_field=col_inds, use_split=True) if bc_temp is not None: col_bcs.append(bc_temp) submat_ctx = ImplicitMatrixContext(asub, row_bcs=row_bcs, col_bcs=col_bcs, fc_params=self.fc_params, appctx=self.appctx) submat_ctx.on_diag = self.on_diag and row_inds == col_inds submat = PETSc.Mat().create(comm=mat.comm) submat.setType("python") submat.setSizes((submat_ctx.row_sizes, submat_ctx.col_sizes), bsize=submat_ctx.block_size) submat.setPythonContext(submat_ctx) submat.setUp() return submat
def plot_matrix_mixed(a_form, bcs=[], **kwargs): """Provides a plot of a mixed matrix.""" fig, ax = plt.subplots(1, 1) A = assemble(a_form, bcs=bcs, mat_type="aij") petsc_mat = A.M.handle total_size = petsc_mat.getSize() f0_size = A.M[0, 0].handle.getSize() Mij = PETSc.Mat() petsc_mat.convert("aij", Mij) n, m = total_size Mnp = np.array(Mij.getValues(range(n), range(m))) Am = np.ma.masked_values(Mnp, 0, rtol=1e-13) # Plot the matrix plot = ax.matshow(Am, cmap=my_cmap, **kwargs) # Remove axis ticks and values ax.tick_params(length=0) ax.set_xticklabels([]) ax.set_yticklabels([]) ax.axhline(y=f0_size[0] - 0.5, color="k") ax.axvline(x=f0_size[0] - 0.5, color="k") return plot
def plot_matrix_hybrid_multiplier_spp(a_form, bcs=[], **kwargs): """Provides a plot of a condensed hybrid-mixed matrix for single scale problems.""" fig, ax = plt.subplots(1, 1) _A = Tensor(a_form) A = _A.blocks S = A[2, 2] - A[2, :2] * A[:2, :2].inv * A[:2, 2] Smat = assemble(S, bcs=bcs) petsc_mat = Smat.M.handle total_size = petsc_mat.getSize() Mij = PETSc.Mat() petsc_mat.convert("aij", Mij) n, m = total_size Mnp = np.array(Mij.getValues(range(n), range(m))) Am = np.ma.masked_values(Mnp, 0, rtol=1e-13) # Plot the matrix plot = ax.matshow(Am, cmap=my_cmap, **kwargs) # Below there is the spy alternative # plot = plt.spy(Am, **kwargs) # Remove axis ticks and values ax.tick_params(length=0) ax.set_xticklabels([]) ax.set_yticklabels([]) return plot
def create_interpolation(dmc, dmf): cctx = firedrake.dmhooks.get_appctx(dmc) fctx = firedrake.dmhooks.get_appctx(dmf) manager = firedrake.dmhooks.get_transfer_manager(dmf) V_c = cctx._problem.u.function_space() V_f = fctx._problem.u.function_space() row_size = V_f.dof_dset.layout_vec.getSizes() col_size = V_c.dof_dset.layout_vec.getSizes() cfn = firedrake.Function(V_c) ffn = firedrake.Function(V_f) cbcs = cctx._problem.bcs fbcs = fctx._problem.bcs ctx = Interpolation(cfn, ffn, manager, cbcs, fbcs) mat = PETSc.Mat().create(comm=dmc.comm) mat.setSizes((row_size, col_size)) mat.setType(mat.Type.PYTHON) mat.setPythonContext(ctx) mat.setUp() return mat, None
def test_ForcingCovariance_init(fs): "test init method of ForcingCovariance" # note: tests only handle case of a single process # need to think about testing more complicated cases # simple example in 1D sigma = np.log(1.) l = np.log(0.1) fc = ForcingCovariance(fs, sigma, l) n = Function(fs).vector().size() n_local = Function(fs).vector().local_size() M = PETSc.Mat().create() M.setSizes(((n_local, -1), (n_local, -1))) M.setFromOptions() M.setUp() start, end = M.getOwnershipRange() assert fc.nx == n assert fc.nx_local == n_local assert fc.function_space == fs assert_allclose(fc.sigma, sigma) assert_allclose(fc.l, l) assert_allclose(fc.cutoff, 1.e-3) assert_allclose(fc.regularization, 1.e-8) assert fc.local_startind == start assert fc.local_endind == end assert not fc.is_assembled
def createSubMatrix(self, mat, row_is, col_is, target=None): if target is not None: # Repeat call, just return the matrix, since we don't # actually assemble in here. target.assemble() return target from firedrake import DirichletBC # These are the sets of ISes of which the the row and column # space consist. row_ises = self._y.function_space().dof_dset.field_ises col_ises = self._x.function_space().dof_dset.field_ises row_inds = find_sub_block(row_is, row_ises) if row_is == col_is and row_ises == col_ises: col_inds = row_inds else: col_inds = find_sub_block(col_is, col_ises) asub = ExtractSubBlock().split(self.a, argument_indices=(row_inds, col_inds)) Wrow = asub.arguments()[0].function_space() Wcol = asub.arguments()[1].function_space() row_bcs = [] col_bcs = [] for bc in self.row_bcs: for i, r in enumerate(row_inds): if bc.function_space().index == r: row_bcs.append(DirichletBC(Wrow.split()[i], bc.function_arg, bc.sub_domain, method=bc.method)) if Wrow == Wcol and row_inds == col_inds and self.row_bcs == self.col_bcs: col_bcs = row_bcs else: for bc in self.col_bcs: for i, c in enumerate(col_inds): if bc.function_space().index == c: col_bcs.append(DirichletBC(Wcol.split()[i], bc.function_arg, bc.sub_domain, method=bc.method)) submat_ctx = ImplicitMatrixContext(asub, row_bcs=row_bcs, col_bcs=col_bcs, fc_params=self.fc_params, appctx=self.appctx) submat_ctx.on_diag = self.on_diag and row_inds == col_inds submat = PETSc.Mat().create(comm=mat.comm) submat.setType("python") submat.setSizes((submat_ctx.row_sizes, submat_ctx.col_sizes), bsize=submat_ctx.block_size) submat.setPythonContext(submat_ctx) submat.setUp() return submat
def get_np_matrix(form, **kwargs): L = get_PETSC_matrix(form, **kwargs) size = L.getSize() Lij = PETSc.Mat() L.convert('aij', Lij) Lnp = np.array(Lij.getValues(range(size[0]), range(size[1]))) return Lnp
def getNestSubMatrix(self, i, j): if i == j: s = self.standalones[i] sizes = (s.Vf.dof_dset.layout_vec.getSizes(), s.Vc.dof_dset.layout_vec.getSizes()) M_shll = PETSc.Mat().createPython(sizes, s, comm=s.Vf.mesh().comm) M_shll.setUp() return M_shll else: return None
def prolongation_matrix_matfree(Vf, Vc, Vf_bcs, Vc_bcs): fele = Vf.ufl_element() if isinstance(fele, MixedElement) and not isinstance(fele, (VectorElement, TensorElement)): ctx = MixedInterpolationMatrix(Vf, Vc, Vf_bcs, Vc_bcs) else: ctx = StandaloneInterpolationMatrix(Vf, Vc, Vf_bcs, Vc_bcs) sizes = (Vc.dof_dset.layout_vec.getSizes(), Vf.dof_dset.layout_vec.getSizes()) M_shll = PETSc.Mat().createPython(sizes, ctx) M_shll.setUp() return M_shll
def GetPETScBFBTApproximateToSchurInverse(): Fass = assemble(lhs(F)) Gass = assemble(inner(u, v) * dx) # bcin.apply(Fass) # bcnoslip.apply(Fass) # bcin.apply(Gass) # bcnoslip.apply(Gass) A = Fass.M[0,0].handle BT = Fass.M[0,1].handle B = Fass.M[1,0].handle D = Fass.M[1,1].handle VM = Gass.M[0, 0].handle # Adiag = A.getDiagonal() Adiag = VM.getRowSum() invAdiag = Adiag.duplicate() Adiag.copy(invAdiag) invAdiag.reciprocal() Einv = A.duplicate() Einv.setDiagonal(invAdiag) # Einv.convert(PETSc.Mat.Type.SEQAIJ) BEBT = B.matMult(Einv.matMult(BT)) BEBT_ksp = PETSc.KSP().create() BEBT_ksp.setOperators(BEBT) opts = PETSc.Options() BEBT_ksp.setOptionsPrefix("bebt_") opts['bebt_ksp_type'] = 'richardson' opts['bebt_pc_type'] = 'hypre' opts['bebt_ksp_max_it'] = 1 opts['bebt_ksp_atol'] = 1.0e-9 opts['bebt_ksp_rtol'] = 1.0e-9 BEBT_ksp.setUp() BEBT_ksp.setFromOptions() MIDDLE = B.matMult(Einv.matMult(A.matMult(Einv.matMult(BT)))) MIDDLE.scale(-1) class SchurInvApprox(object): def mult(self, mat, x, y): y1 = y.duplicate() BEBT_ksp.solve(x, y1) y2 = y.duplicate() MIDDLE.mult(y1, y2) BEBT_ksp.solve(y2, y) schur = PETSc.Mat() schur.createPython(D.getSizes(), SchurInvApprox()) schur.setUp() return schur
def create_transfer(cctx, fctx, mat_type, cbcs, fbcs): cV = cctx.J.arguments()[0].function_space() fV = fctx.J.arguments()[0].function_space() cbcs = cctx._problem.bcs if cbcs else [] fbcs = fctx._problem.bcs if fbcs else [] if mat_type == "matfree": I = prolongation_matrix_matfree(fV, cV, fbcs, cbcs) elif mat_type == "aij": I = PETSc.Mat().createTranspose(prolongation_matrix_aij(fV, cV, fbcs, cbcs)) else: raise ValueError("Unknown matrix type") return I
def GetPETScInverseMassApproximateToSchurInverse(): mass = assemble( p*q*dx).M[1,1].handle mass_diag = mass.createVecLeft() mass.getDiagonal(mass_diag) mass_diag.reciprocal() mass_diag.scale(-1) class InvMassSchurInvApprox(object): def mult(self, mat, x, y): y.pointwiseMult(mass_diag, x) schur = PETSc.Mat() schur.createPython(mass.getSizes(), InvMassSchurInvApprox()) schur.setUp() return schur
def duplicate(self, mat, copy): if copy == 0: raise NotImplementedError("We do now know how to duplicate a matrix-free MAT when copy=0") newmat_ctx = ImplicitMatrixContext(self.a, row_bcs=self.bcs, col_bcs=self.bcs_col, fc_params=self.fc_params, appctx=self.appctx) newmat = PETSc.Mat().create(comm=mat.comm) newmat.setType("python") newmat.setSizes((newmat_ctx.row_sizes, newmat_ctx.col_sizes), bsize=newmat_ctx.block_size) newmat.setPythonContext(newmat_ctx) newmat.setUp() return newmat
def create_interpolation(self, dmc, dmf): # This should be generalised to work for arbitrary function # spaces. Currently I think it only works for CG/DG on simplices. # I used the same code as firedrake.P1PC. cctx = get_appctx(dmc) fctx = get_appctx(dmf) cV = cctx.J.arguments()[0].function_space() fV = fctx.J.arguments()[0].function_space() cbcs = cctx._problem.bcs fbcs = fctx._problem.bcs I = prolongation_matrix(fV, cV, fbcs, cbcs) R = PETSc.Mat().createTranspose(I) return R, None
def __init__(self, function_space, coords): "create and assemble interpolation matrix" if not isinstance(function_space, WithGeometry): raise TypeError( "bad input type for function_space: must be a FunctionSpace") self.coords = np.copy(coords) self.function_space = function_space self.comm = function_space.comm self.n_data = coords.shape[0] assert (coords.shape[1] == self.function_space.mesh().cell_dimension() ), "shape of coordinates does not match mesh dimension" # allocate working vectors to handle parallel matrix operations and data transfer # dataspace_vector is a distributed PETSc vector in the data space self.dataspace_distrib = PETSc.Vec().create(comm=self.comm) self.dataspace_distrib.setSizes((-1, self.n_data)) self.dataspace_distrib.setFromOptions() self.n_data_local = self.dataspace_distrib.getSizes()[0] # all data computations are done on root process, so create gather method to # facilitate this data transfer self.petsc_scatter, self.dataspace_gathered = PETSc.Scatter.toZero( self.dataspace_distrib) self.meshspace_vector = Function(self.function_space).vector() self.n_mesh_local = self.meshspace_vector.local_size() self.n_mesh = self.meshspace_vector.size() nnz = len(self.function_space.cell_node_list[0]) self.interp = PETSc.Mat().create(comm=self.comm) self.interp.setSizes( ((self.n_mesh_local, -1), (self.n_data_local, -1))) self.interp.setPreallocationNNZ(nnz) self.interp.setFromOptions() self.interp.setUp() self.is_assembled = False
def construct_full_interpolation_matrix(self, IFW): """ Assemble interpolation matrix for vectorial tensorized spline space. """ # this is one of the two bottlenecks that slow down initiating Bsplines FullIFW = PETSc.Mat().create(self.comm) FullIFW.setType(PETSc.Mat.Type.AIJ) # set proper matrix sizes d = self.dim free_dims = list(set(range(self.dim)) - set(self.fixed_dims)) dfree = len(free_dims) ((lsize, gsize), (lsize_spline, gsize_spline)) = IFW.getSizes() FullIFW.setSizes(((d * lsize, d * gsize), (dfree * lsize_spline, dfree * gsize_spline))) # (over)estimate sparsity pattern using row with most nonzeros # possible memory improvement: allocate precise sparsity pattern # row by row (but this needs nnzdiagonal and nnzoffidagonal; # not straightforward to do) global_rows = self.lg_map_fe.apply([range(lsize)]) for ii, row in enumerate(range(lsize)): row = global_rows[row] self.FullIFWnnz = max(self.FullIFWnnz, len(IFW.getRow(row)[1])) FullIFW.setPreallocationNNZ(self.FullIFWnnz) # preallocate matrix FullIFW.setUp() # fill matrix by blowing up entries from IFW to do the right thing # on vector fields (it's not just a block matrix: values are # interleaved as this is how firedrake handles vector fields) innerloop_idx = [[i, free_dims[i]] for i in range(dfree)] for row in range(lsize): # for every FE dof row = self.lg_map_fe.apply([row])[0] # extract value of all tensorize Bsplines at this dof (cols, vals) = IFW.getRow(row) expandedcols = dfree * cols for j, dim in innerloop_idx: FullIFW.setValues( [d * row + dim], # global row expandedcols + j, # global column vals) FullIFW.assemble() return FullIFW
def solve(ctx, state): out_file = File(solution_out) if solution_out else None mass = state["mass"] hats = state["hats"] u = ctx[2]["u"] u0 = ctx[2]["u0"] solver = ctx[1] from firedrake.petsc import PETSc ksp_hats = PETSc.KSP() ksp_hats.create() ksp_hats.setOperators(hats) opts = PETSc.Options() opts['ksp_type'] = inner_ksp opts['ksp_max_it'] = max_iterations opts['pc_type'] = 'hypre' ksp_hats.setFromOptions() class SchurInv(object): def mult(self, mat, x, y): tmp1 = y.duplicate() tmp2 = y.duplicate() ksp_hats.solve(x, tmp1) mass.mult(tmp1, tmp2) ksp_hats.solve(tmp2, y) pc_schur = PETSc.Mat() pc_schur.createPython(mass.getSizes(), SchurInv()) pc_schur.setUp() pc = solver.snes.ksp.pc pc.setFieldSplitSchurPreType(PETSc.PC.SchurPreType.USER, pc_schur) for step in range(steps): casper.invoke_task(ctx, "assign", state) solver.solve() if out_file is not None: out_file.write(u.split()[0], time=step) if compute_norms: nu = norm(u) if comm.rank == 0: print(step, 'L2(u):', nu)
def __init__(self, a, bcs, *args, **kwargs): # sets self._a and self._bcs super(ImplicitMatrix, self).__init__(a, bcs) appctx = kwargs.get("appctx", {}) from firedrake.matrix_free.operators import ImplicitMatrixContext ctx = ImplicitMatrixContext(a, row_bcs=self.bcs, col_bcs=self.bcs, fc_params=kwargs["fc_params"], appctx=appctx) self.petscmat = PETSc.Mat().create(comm=self.comm) self.petscmat.setType("python") self.petscmat.setSizes((ctx.row_sizes, ctx.col_sizes), bsize=ctx.block_size) self.petscmat.setPythonContext(ctx) self.petscmat.setUp() self.petscmat.assemble()
def create_interpolation(self, dmc, dmf): cctx = get_appctx(dmc) fctx = get_appctx(dmf) cV = cctx.J.arguments()[0].function_space() fV = fctx.J.arguments()[0].function_space() cbcs = cctx._problem.bcs fbcs = fctx._problem.bcs prefix = dmc.getOptionsPrefix() mattype = PETSc.Options(prefix).getString("mg_levels_transfer_mat_type", default="matfree") if mattype == "matfree": I = prolongation_matrix_matfree(fV, cV, fbcs, cbcs) elif mattype == "aij": I = prolongation_matrix_aij(fV, cV, fbcs, cbcs) else: raise ValueError("Unknown matrix type") R = PETSc.Mat().createTranspose(I) return R, None
def create_injection(dmc, dmf): _, clvl = utils.get_level(dmc) _, flvl = utils.get_level(dmf) cctx = dmc.getAppCtx() V_c = dmc.getAttr("__fs__")() V_f = dmf.getAttr("__fs__")() nrow = sum(x.dof_dset.size * x.dof_dset.cdim for x in V_f) ncol = sum(x.dof_dset.size * x.dof_dset.cdim for x in V_c) cfn = firedrake.Function(V_c) ffn = firedrake.Function(V_f) cbcs = cctx._problems[clvl].bcs class Injection(object): def __init__(self, cfn, ffn, cbcs=None): self.cfn = cfn self.ffn = ffn self.cbcs = cbcs or [] def multTranspose(self, mat, x, y): with self.ffn.dat.vec as v: x.copy(v) firedrake.inject(self.ffn, self.cfn) for bc in self.cbcs: bc.apply(self.cfn) with self.cfn.dat.vec_ro as v: v.copy(y) ctx = Injection(cfn, ffn, cbcs) mat = PETSc.Mat().create() mat.setSizes(((nrow, None), (ncol, None))) mat.setType(mat.Type.PYTHON) mat.setPythonContext(ctx) mat.setUp() return mat
def construct_kronecker_matrix(self, interp_1d): """ Construct the tensorized interpolation matrix. Do this by computing the kron product of the rows of the 1d univariate interpolation matrices. In the future, this may be done matrix-free. """ # this is one of the two bottlenecks that slow down initiating Bsplines IFW = PETSc.Mat().create(self.comm) IFW.setType(PETSc.Mat.Type.AIJ) comm = self.comm # owned part of global problem local_N = self.N // comm.size + int(comm.rank < (self.N % comm.size)) (lsize, gsize) = interp_1d[0].getSizes()[0] IFW.setSizes(((lsize, gsize), (local_N, self.N))) # guess sparsity pattern from interp_1d[0] for row in range(lsize): row = self.lg_map_fe.apply([row])[0] nnz_ = len(interp_1d[0].getRow(row)[0]) # lenght of nnz-array self.IFWnnz = max(self.IFWnnz, nnz_**self.dim) IFW.setPreallocationNNZ(self.IFWnnz) IFW.setUp() for row in range(lsize): row = self.lg_map_fe.apply([row])[0] M = [[A.getRow(row)[0], A.getRow(row)[1], A.getSize()[1]] for A in interp_1d] M = reduce(self.vectorkron, M) columns, values, lenght = M IFW.setValues([row], columns, values) IFW.assemble() return IFW
def nonlocal_integral_eq( mesh, scatterer_bdy_id, outer_bdy_id, wave_number, options_prefix=None, solver_parameters=None, fspace=None, vfspace=None, true_sol_grad_expr=None, actx=None, dgfspace=None, dgvfspace=None, meshmode_src_connection=None, qbx_kwargs=None, ): r""" see run_method for descriptions of unlisted args args: gamma and beta are used to precondition with the following equation: \Delta u - \kappa^2 \gamma u = 0 (\partial_n - i\kappa\beta) u |_\Sigma = 0 """ # make sure we get outer bdy id as tuple in case it consists of multiple ids if isinstance(outer_bdy_id, int): outer_bdy_id = [outer_bdy_id] outer_bdy_id = tuple(outer_bdy_id) # away from the excluded region, but firedrake and meshmode point # into pyt_inner_normal_sign = -1 ambient_dim = mesh.geometric_dimension() # {{{ Build src and tgt # build connection meshmode near src boundary -> src boundary inside meshmode from meshmode.discretization.poly_element import \ InterpolatoryQuadratureSimplexGroupFactory from meshmode.discretization.connection import make_face_restriction factory = InterpolatoryQuadratureSimplexGroupFactory( dgfspace.finat_element.degree) src_bdy_connection = make_face_restriction(actx, meshmode_src_connection.discr, factory, scatterer_bdy_id) # source is a qbx layer potential from pytential.qbx import QBXLayerPotentialSource disable_refinement = (fspace.mesh().geometric_dimension() == 3) qbx = QBXLayerPotentialSource(src_bdy_connection.to_discr, **qbx_kwargs, _disable_refinement=disable_refinement) # get target indices and point-set target_indices, target = get_target_points_and_indices( fspace, outer_bdy_id) # }}} # build the operations from pytential import bind, sym r""" ..math: x \in \Sigma grad_op(x) = \nabla( \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ) """ grad_op = pyt_inner_normal_sign * sym.grad( ambient_dim, sym.D(HelmholtzKernel(ambient_dim), sym.var("u"), k=sym.var("k"), qbx_forced_limit=None)) r""" ..math: x \in \Sigma op(x) = i \kappa \cdot \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y) """ op = pyt_inner_normal_sign * 1j * sym.var("k") * (sym.D( HelmholtzKernel(ambient_dim), sym.var("u"), k=sym.var("k"), qbx_forced_limit=None)) # bind the operations pyt_grad_op = bind((qbx, target), grad_op) pyt_op = bind((qbx, target), op) # }}} class MatrixFreeB(object): def __init__(self, A, pyt_grad_op, pyt_op, actx, kappa): """ :arg kappa: The wave number """ self.actx = actx self.k = kappa self.pyt_op = pyt_op self.pyt_grad_op = pyt_grad_op self.A = A self.meshmode_src_connection = meshmode_src_connection # {{{ Create some functions needed for multing self.x_fntn = Function(fspace) # CG self.potential_int = Function(fspace) self.potential_int.dat.data[:] = 0.0 self.grad_potential_int = Function(vfspace) self.grad_potential_int.dat.data[:] = 0.0 self.pyt_result = Function(fspace) self.n = FacetNormal(mesh) self.v = TestFunction(fspace) # some meshmode ones self.x_mm_fntn = self.meshmode_src_connection.discr.empty( self.actx, dtype='c') # }}} def mult(self, mat, x, y): # Copy function data into the fivredrake function self.x_fntn.dat.data[:] = x[:] # Transfer the function to meshmode self.meshmode_src_connection.from_firedrake(project( self.x_fntn, dgfspace), out=self.x_mm_fntn) # Restrict to boundary x_mm_fntn_on_bdy = src_bdy_connection(self.x_mm_fntn) # Apply the operation potential_int_mm = self.pyt_op(self.actx, u=x_mm_fntn_on_bdy, k=self.k) grad_potential_int_mm = self.pyt_grad_op(self.actx, u=x_mm_fntn_on_bdy, k=self.k) # Store in firedrake self.potential_int.dat.data[target_indices] = potential_int_mm.get( ) for dim in range(grad_potential_int_mm.shape[0]): self.grad_potential_int.dat.data[ target_indices, dim] = grad_potential_int_mm[dim].get() # Integrate the potential r""" Compute the inner products using firedrake. Note this will be subtracted later, hence appears off by a sign. .. math:: \langle n(x) \cdot \nabla( \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ), v \rangle_\Sigma - \langle i \kappa \cdot \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y), v \rangle_\Sigma """ self.pyt_result = assemble( inner(inner(self.grad_potential_int, self.n), self.v) * ds(outer_bdy_id) - inner(self.potential_int, self.v) * ds(outer_bdy_id)) # y <- Ax - evaluated potential self.A.mult(x, y) with self.pyt_result.dat.vec_ro as ep: y.axpy(-1, ep) # {{{ Compute normal helmholtz operator u = TrialFunction(fspace) v = TestFunction(fspace) r""" .. math:: \langle \nabla u, \nabla v \rangle - \kappa^2 \cdot \langle u, v \rangle - i \kappa \langle u, v \rangle_\Sigma """ a = inner(grad(u), grad(v)) * dx \ - Constant(wave_number**2) * inner(u, v) * dx \ - Constant(1j * wave_number) * inner(u, v) * ds(outer_bdy_id) # get the concrete matrix from a general bilinear form A = assemble(a).M.handle # }}} # {{{ Setup Python matrix B = PETSc.Mat().create() # build matrix context Bctx = MatrixFreeB(A, pyt_grad_op, pyt_op, actx, wave_number) # set up B as same size as A B.setSizes(*A.getSizes()) B.setType(B.Type.PYTHON) B.setPythonContext(Bctx) B.setUp() # }}} # {{{ Create rhs # Remember f is \partial_n(true_sol)|_\Gamma # so we just need to compute \int_\Gamma\partial_n(true_sol) H(x-y) sigma = sym.make_sym_vector("sigma", ambient_dim) r""" ..math: x \in \Sigma grad_op(x) = \nabla( \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ) """ grad_op = pyt_inner_normal_sign * \ sym.grad(ambient_dim, sym.S(HelmholtzKernel(ambient_dim), sym.n_dot(sigma), k=sym.var("k"), qbx_forced_limit=None)) r""" ..math: x \in \Sigma op(x) = i \kappa \cdot \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ) """ op = 1j * sym.var("k") * pyt_inner_normal_sign * \ sym.S(HelmholtzKernel(ambient_dim), sym.n_dot(sigma), k=sym.var("k"), qbx_forced_limit=None) rhs_grad_op = bind((qbx, target), grad_op) rhs_op = bind((qbx, target), op) # Transfer to meshmode metadata = {'quadrature_degree': 2 * fspace.ufl_element().degree()} dg_true_sol_grad = project(true_sol_grad_expr, dgvfspace, form_compiler_parameters=metadata) true_sol_grad_mm = meshmode_src_connection.from_firedrake(dg_true_sol_grad, actx=actx) true_sol_grad_mm = src_bdy_connection(true_sol_grad_mm) # Apply the operations f_grad_convoluted_mm = rhs_grad_op(actx, sigma=true_sol_grad_mm, k=wave_number) f_convoluted_mm = rhs_op(actx, sigma=true_sol_grad_mm, k=wave_number) # Transfer function back to firedrake f_grad_convoluted = Function(vfspace) f_convoluted = Function(fspace) f_grad_convoluted.dat.data[:] = 0.0 f_convoluted.dat.data[:] = 0.0 for dim in range(f_grad_convoluted_mm.shape[0]): f_grad_convoluted.dat.data[target_indices, dim] = f_grad_convoluted_mm[dim].get() f_convoluted.dat.data[target_indices] = f_convoluted_mm.get() r""" \langle f, v \rangle_\Gamma + \langle i \kappa \cdot \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y), v \rangle_\Sigma - \langle n(x) \cdot \nabla( \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ), v \rangle_\Sigma """ rhs_form = inner(inner(true_sol_grad_expr, FacetNormal(mesh)), v) * ds(scatterer_bdy_id, metadata=metadata) \ + inner(f_convoluted, v) * ds(outer_bdy_id) \ - inner(inner(f_grad_convoluted, FacetNormal(mesh)), v) * ds(outer_bdy_id) rhs = assemble(rhs_form) # {{{ set up a solver: solution = Function(fspace, name="Computed Solution") # {{{ Used for preconditioning if 'gamma' in solver_parameters or 'beta' in solver_parameters: gamma = complex(solver_parameters.pop('gamma', 1.0)) import cmath beta = complex(solver_parameters.pop('beta', cmath.sqrt(gamma))) p = inner(grad(u), grad(v)) * dx \ - Constant(wave_number**2 * gamma) * inner(u, v) * dx \ - Constant(1j * wave_number * beta) * inner(u, v) * ds(outer_bdy_id) P = assemble(p).M.handle else: P = A # }}} # Set up options to contain solver parameters: ksp = PETSc.KSP().create() if solver_parameters['pc_type'] == 'pyamg': del solver_parameters['pc_type'] # We are using the AMG preconditioner pyamg_tol = solver_parameters.get('pyamg_tol', None) if pyamg_tol is not None: pyamg_tol = float(pyamg_tol) pyamg_maxiter = solver_parameters.get('pyamg_maxiter', None) if pyamg_maxiter is not None: pyamg_maxiter = int(pyamg_maxiter) ksp.setOperators(B) ksp.setUp() pc = ksp.pc pc.setType(pc.Type.PYTHON) pc.setPythonContext( AMGTransmissionPreconditioner(wave_number, fspace, A, tol=pyamg_tol, maxiter=pyamg_maxiter, use_plane_waves=True)) # Otherwise use regular preconditioner else: ksp.setOperators(B, P) options_manager = OptionsManager(solver_parameters, options_prefix) options_manager.set_from_options(ksp) import petsc4py.PETSc petsc4py.PETSc.Sys.popErrorHandler() with rhs.dat.vec_ro as b: with solution.dat.vec as x: ksp.solve(b, x) # }}} return ksp, solution
def assemble_mixed_mass_matrix(V_A, V_B): """ Construct the mixed mass matrix of two function spaces, using the TrialFunction from V_A and the TestFunction from V_B. """ if len(V_A) > 1 or len(V_B) > 1: raise NotImplementedError( "Sorry, only implemented for non-mixed spaces") if V_A.ufl_element().mapping() != "identity" or V_B.ufl_element().mapping( ) != "identity": msg = """ Sorry, only implemented for affine maps for now. To do non-affine, we'd need to import much more of the assembly engine of UFL/TSFC/etc to do the assembly on each supermesh cell. """ raise NotImplementedError(msg) mesh_A = V_A.mesh() mesh_B = V_B.mesh() dim = mesh_A.geometric_dimension() assert dim == mesh_B.geometric_dimension() assert dim == mesh_A.topological_dimension() assert dim == mesh_B.topological_dimension() (mh_A, level_A) = get_level(mesh_A) (mh_B, level_B) = get_level(mesh_B) if mesh_A is mesh_B: def likely(cell_A): return [cell_A] else: if (mh_A is None or mh_B is None) or (mh_A is not mh_B): # No mesh hierarchy structure, call libsupermesh for # intersection finding intersections = intersection_finder(mesh_A, mesh_B) likely = intersections.__getitem__ else: # We do have a mesh hierarchy, use it if abs(level_A - level_B) > 1: raise NotImplementedError( "Only works for transferring between adjacent levels for now." ) # What are the cells of B that (probably) intersect with a given cell in A? if level_A > level_B: cell_map = mh_A.fine_to_coarse_cells[level_A] def likely(cell_A): return cell_map[cell_A] elif level_A < level_B: cell_map = mh_A.coarse_to_fine_cells[level_A] def likely(cell_A): return cell_map[cell_A] assert V_A.value_size == V_B.value_size orig_value_size = V_A.value_size if V_A.value_size > 1: V_A = firedrake.FunctionSpace(mesh_A, V_A.ufl_element().sub_elements()[0]) if V_B.value_size > 1: V_B = firedrake.FunctionSpace(mesh_B, V_B.ufl_element().sub_elements()[0]) assert V_A.value_size == 1 assert V_B.value_size == 1 preallocator = PETSc.Mat().create(comm=mesh_A.comm) preallocator.setType(PETSc.Mat.Type.PREALLOCATOR) rset = V_B.dof_dset cset = V_A.dof_dset nrows = rset.layout_vec.getSizes() ncols = cset.layout_vec.getSizes() preallocator.setLGMap(rmap=rset.scalar_lgmap, cmap=cset.scalar_lgmap) preallocator.setSizes(size=(nrows, ncols), bsize=1) preallocator.setUp() zeros = numpy.zeros((V_B.cell_node_map().arity, V_A.cell_node_map().arity), dtype=ScalarType) for cell_A, dofs_A in enumerate(V_A.cell_node_map().values): for cell_B in likely(cell_A): dofs_B = V_B.cell_node_map().values_with_halo[cell_B, :] preallocator.setValuesLocal(dofs_B, dofs_A, zeros) preallocator.assemble() dnnz, onnz = get_preallocation(preallocator, nrows[0]) # Unroll from block to AIJ dnnz = dnnz * cset.cdim dnnz = numpy.repeat(dnnz, rset.cdim) onnz = onnz * cset.cdim onnz = numpy.repeat(onnz, cset.cdim) preallocator.destroy() assert V_A.value_size == V_B.value_size rdim = V_B.dof_dset.cdim cdim = V_A.dof_dset.cdim # # Preallocate M_AB. # mat = PETSc.Mat().create(comm=mesh_A.comm) mat.setType(PETSc.Mat.Type.AIJ) rsizes = tuple(n * rdim for n in nrows) csizes = tuple(c * cdim for c in ncols) mat.setSizes(size=(rsizes, csizes), bsize=(rdim, cdim)) mat.setPreallocationNNZ((dnnz, onnz)) mat.setLGMap(rmap=rset.lgmap, cmap=cset.lgmap) # TODO: Boundary conditions not handled. mat.setOption(mat.Option.IGNORE_OFF_PROC_ENTRIES, False) mat.setOption(mat.Option.NEW_NONZERO_ALLOCATION_ERR, True) mat.setOption(mat.Option.KEEP_NONZERO_PATTERN, True) mat.setOption(mat.Option.UNUSED_NONZERO_LOCATION_ERR, False) mat.setOption(mat.Option.IGNORE_ZERO_ENTRIES, True) mat.setUp() evaluate_kernel_A = compile_element(ufl.Coefficient(V_A), name="evaluate_kernel_A") evaluate_kernel_B = compile_element(ufl.Coefficient(V_B), name="evaluate_kernel_B") # We only need one of these since we assume that the two meshes both have CG1 coordinates to_reference_kernel = to_reference_coordinates( mesh_A.coordinates.ufl_element()) if dim == 2: reference_mesh = UnitTriangleMesh(comm=COMM_SELF) else: reference_mesh = UnitTetrahedronMesh(comm=COMM_SELF) evaluate_kernel_S = compile_element(ufl.Coefficient( reference_mesh.coordinates.function_space()), name="evaluate_kernel_S") V_S_A = FunctionSpace(reference_mesh, V_A.ufl_element()) V_S_B = FunctionSpace(reference_mesh, V_B.ufl_element()) M_SS = assemble(inner(TrialFunction(V_S_A), TestFunction(V_S_B)) * dx) M_SS = M_SS.M.handle[:, :] node_locations_A = utils.physical_node_locations( V_S_A).dat.data_ro_with_halos node_locations_B = utils.physical_node_locations( V_S_B).dat.data_ro_with_halos num_nodes_A = node_locations_A.shape[0] num_nodes_B = node_locations_B.shape[0] to_reference_kernel = to_reference_coordinates( mesh_A.coordinates.ufl_element()) supermesh_kernel_str = """ #include "libsupermesh-c.h" #include <petsc.h> %(to_reference)s %(evaluate_S)s %(evaluate_A)s %(evaluate_B)s #define complex_mode %(complex_mode)s #define PrintInfo(...) do { if (PetscLogPrintInfo) printf(__VA_ARGS__); } while (0) static void print_array(PetscScalar *arr, int d) { for(int j=0; j<d; j++) PrintInfo(stderr, "%%+.2f ", arr[j]); } static void print_coordinates(PetscScalar *simplex, int d) { for(int i=0; i<d+1; i++) { PrintInfo("\t"); print_array(&simplex[d*i], d); PrintInfo("\\n"); } } #if complex_mode static void seperate_real_and_imag(PetscScalar *simplex, double *real_simplex, double *imag_simplex, int d) { for(int i=0; i<d+1; i++) { for(int j=0; j<d; j++) { real_simplex[d*i+j] = creal(simplex[d*i+j]); imag_simplex[d*i+j] = cimag(simplex[d*i+j]); } } } static void merge_back_to_simplex(PetscScalar* simplex, double* real_simplex, double* imag_simplex, int d) { print_coordinates(simplex,d); for(int i=0; i<d+1; i++) { for(int j=0; j<d; j++) { simplex[d*i+j] = real_simplex[d*i+j]+imag_simplex[d*i+j]*_Complex_I; } } } #endif int supermesh_kernel(PetscScalar* simplex_A, PetscScalar* simplex_B, PetscScalar* simplices_C, PetscScalar* nodes_A, PetscScalar* nodes_B, PetscScalar* M_SS, PetscScalar* outptr, int num_ele) { #define d %(dim)s #define num_nodes_A %(num_nodes_A)s #define num_nodes_B %(num_nodes_B)s double simplex_ref_measure; PrintInfo("simplex_A coordinates\\n"); print_coordinates(simplex_A, d); PrintInfo("simplex_B coordinates\\n"); print_coordinates(simplex_B, d); int num_elements = num_ele; if (d == 2) simplex_ref_measure = 0.5; else if (d == 3) simplex_ref_measure = 1.0/6; PetscScalar R_AS[num_nodes_A][num_nodes_A]; PetscScalar R_BS[num_nodes_B][num_nodes_B]; PetscScalar coeffs_A[%(num_nodes_A)s] = {0.}; PetscScalar coeffs_B[%(num_nodes_B)s] = {0.}; PetscScalar reference_nodes_A[num_nodes_A][d]; PetscScalar reference_nodes_B[num_nodes_B][d]; #if complex_mode double real_simplex_A[d*(d+1)]; double imag_simplex_A[d*(d+1)]; seperate_real_and_imag(simplex_A, real_simplex_A, imag_simplex_A, d); double real_simplex_B[d*(d+1)]; double imag_simplex_B[d*(d+1)]; seperate_real_and_imag(simplex_B, real_simplex_B, imag_simplex_B, d); double real_simplices_C[num_elements*d*(d+1)]; double imag_simplices_C[num_elements*d*(d+1)]; for (int ii=0; ii<num_elements*d*(d+1); ++ii) imag_simplices_C[ii] = 0.; %(libsupermesh_intersect_simplices)s(real_simplex_A, real_simplex_B, real_simplices_C, &num_elements); merge_back_to_simplex(simplex_A, real_simplex_A, imag_simplex_A, d); merge_back_to_simplex(simplex_B, real_simplex_B, imag_simplex_B, d); for(int s=0; s<num_elements; s++) { PetscScalar* simplex_C = &simplices_C[s * d * (d+1)]; double* real_simplex_C = &real_simplices_C[s * d * (d+1)]; double* imag_simplex_C = &imag_simplices_C[s * d * (d+1)]; merge_back_to_simplex(simplex_C, real_simplex_C, imag_simplex_C, d); } #else %(libsupermesh_intersect_simplices)s(simplex_A, simplex_B, simplices_C, &num_elements); #endif PrintInfo("Supermesh consists of %%i elements\\n", num_elements); // would like to do this //PetscScalar MAB[%(num_nodes_A)s][%(num_nodes_B)s] = (PetscScalar (*)[%(num_nodes_B)s])outptr; // but have to do this instead because we don't grok C PetscScalar (*MAB)[num_nodes_A] = (PetscScalar (*)[num_nodes_A])outptr; PetscScalar (*MSS)[num_nodes_A] = (PetscScalar (*)[num_nodes_A])M_SS; // note the underscore for ( int i = 0; i < num_nodes_B; i++ ) { for (int j = 0; j < num_nodes_A; j++) { MAB[i][j] = 0.0; } } for(int s=0; s<num_elements; s++) { PetscScalar* simplex_S = &simplices_C[s * d * (d+1)]; double simplex_S_measure; #if complex_mode double real_simplex_S[d*(d+1)]; double imag_simplex_S[d*(d+1)]; seperate_real_and_imag(simplex_S, real_simplex_S, imag_simplex_S, d); %(libsupermesh_simplex_measure)s(real_simplex_S, &simplex_S_measure); merge_back_to_simplex(simplex_S, real_simplex_S, imag_simplex_S, d); #else %(libsupermesh_simplex_measure)s(simplex_S, &simplex_S_measure); #endif PrintInfo("simplex_S coordinates with measure %%f\\n", simplex_S_measure); print_coordinates(simplex_S, d); PrintInfo("Start mapping nodes for V_A\\n"); PetscScalar physical_nodes_A[num_nodes_A][d]; for(int n=0; n < num_nodes_A; n++) { PetscScalar* reference_node_location = &nodes_A[n*d]; PetscScalar* physical_node_location = physical_nodes_A[n]; for (int j=0; j < d; j++) physical_node_location[j] = 0.0; pyop2_kernel_evaluate_kernel_S(physical_node_location, simplex_S, reference_node_location); PrintInfo("\\tNode "); print_array(reference_node_location, d); PrintInfo(" mapped to "); print_array(physical_node_location, d); PrintInfo("\\n"); } PrintInfo("Start mapping nodes for V_B\\n"); PetscScalar physical_nodes_B[num_nodes_B][d]; for(int n=0; n < num_nodes_B; n++) { PetscScalar* reference_node_location = &nodes_B[n*d]; PetscScalar* physical_node_location = physical_nodes_B[n]; for (int j=0; j < d; j++) physical_node_location[j] = 0.0; pyop2_kernel_evaluate_kernel_S(physical_node_location, simplex_S, reference_node_location); PrintInfo("\\tNode "); print_array(reference_node_location, d); PrintInfo(" mapped to "); print_array(physical_node_location, d); PrintInfo("\\n"); } PrintInfo("==========================================================\\n"); PrintInfo("Start pulling back dof from S into reference space for A.\\n"); for(int n=0; n < num_nodes_A; n++) { for(int i=0; i<d; i++) reference_nodes_A[n][i] = 0.; to_reference_coords_kernel(reference_nodes_A[n], physical_nodes_A[n], simplex_A); PrintInfo("Pulling back "); print_array(physical_nodes_A[n], d); PrintInfo(" to "); print_array(reference_nodes_A[n], d); PrintInfo("\\n"); } PrintInfo("Start pulling back dof from S into reference space for B.\\n"); for(int n=0; n < num_nodes_B; n++) { for(int i=0; i<d; i++) reference_nodes_B[n][i] = 0.; to_reference_coords_kernel(reference_nodes_B[n], physical_nodes_B[n], simplex_B); PrintInfo("Pulling back "); print_array(physical_nodes_B[n], d); PrintInfo(" to "); print_array(reference_nodes_B[n], d); PrintInfo("\\n"); } PrintInfo("Start evaluating basis functions of V_A at dofs for V_A on S\\n"); for(int i=0; i<num_nodes_A; i++) { coeffs_A[i] = 1.; for(int j=0; j<num_nodes_A; j++) { R_AS[i][j] = 0.; pyop2_kernel_evaluate_kernel_A(&R_AS[i][j], coeffs_A, reference_nodes_A[j]); } print_array(R_AS[i], num_nodes_A); PrintInfo("\\n"); coeffs_A[i] = 0.; } PrintInfo("Start evaluating basis functions of V_B at dofs for V_B on S\\n"); for(int i=0; i<num_nodes_B; i++) { coeffs_B[i] = 1.; for(int j=0; j<num_nodes_B; j++) { R_BS[i][j] = 0.; pyop2_kernel_evaluate_kernel_B(&R_BS[i][j], coeffs_B, reference_nodes_B[j]); } print_array(R_BS[i], num_nodes_B); PrintInfo("\\n"); coeffs_B[i] = 0.; } PrintInfo("Start doing the matmatmat mult\\n"); for ( int i = 0; i < num_nodes_B; i++ ) { for (int j = 0; j < num_nodes_A; j++) { for ( int k = 0; k < num_nodes_B; k++) { for ( int l = 0; l < num_nodes_A; l++) { MAB[i][j] += (simplex_S_measure/simplex_ref_measure) * R_BS[i][k] * MSS[k][l] * R_AS[j][l]; } } } } } return num_elements; } """ % { "evaluate_S": str(evaluate_kernel_S), "evaluate_A": str(evaluate_kernel_A), "evaluate_B": str(evaluate_kernel_B), "to_reference": str(to_reference_kernel), "num_nodes_A": num_nodes_A, "num_nodes_B": num_nodes_B, "libsupermesh_simplex_measure": "libsupermesh_triangle_area" if dim == 2 else "libsupermesh_tetrahedron_volume", "libsupermesh_intersect_simplices": "libsupermesh_intersect_tris_real" if dim == 2 else "libsupermesh_intersect_tets_real", "dim": dim, "complex_mode": 1 if complex_mode else 0 } dirs = get_petsc_dir() + (sys.prefix, ) includes = ["-I%s/include" % d for d in dirs] libs = ["-L%s/lib" % d for d in dirs] libs = libs + ["-Wl,-rpath,%s/lib" % d for d in dirs] + ["-lpetsc", "-lsupermesh"] lib = load(supermesh_kernel_str, "c", "supermesh_kernel", cppargs=includes, ldargs=libs, argtypes=[ ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp ], restype=ctypes.c_int) ammm(V_A, V_B, likely, node_locations_A, node_locations_B, M_SS, ctypes.addressof(lib), mat) if orig_value_size == 1: return mat else: (lrows, grows), (lcols, gcols) = mat.getSizes() lrows *= orig_value_size grows *= orig_value_size lcols *= orig_value_size gcols *= orig_value_size size = ((lrows, grows), (lcols, gcols)) context = BlockMatrix(mat, orig_value_size) blockmat = PETSc.Mat().createPython(size, context=context, comm=mat.comm) blockmat.setUp() return blockmat
def create_interpolation(dmc, dmf): _, clvl = utils.get_level(dmc) _, flvl = utils.get_level(dmf) cctx = dmc.getAppCtx() fctx = dmf.getAppCtx() V_c = dmc.getAttr("__fs__")() V_f = dmf.getAttr("__fs__")() nrow = sum(x.dof_dset.size * x.dof_dset.cdim for x in V_f) ncol = sum(x.dof_dset.size * x.dof_dset.cdim for x in V_c) cfn = firedrake.Function(V_c) ffn = firedrake.Function(V_f) cbcs = cctx._problems[clvl].bcs fbcs = fctx._problems[flvl].bcs class Interpolation(object): def __init__(self, cfn, ffn, cbcs=None, fbcs=None): self.cfn = cfn self.ffn = ffn self.cbcs = cbcs or [] self.fbcs = fbcs or [] def mult(self, mat, x, y, inc=False): with self.cfn.dat.vec as v: x.copy(v) firedrake.prolong(self.cfn, self.ffn) for bc in self.fbcs: bc.zero(self.ffn) with self.ffn.dat.vec_ro as v: if inc: y.axpy(1.0, v) else: v.copy(y) def multAdd(self, mat, x, y, w): if y.handle == w.handle: self.mult(mat, x, w, inc=True) else: self.mult(mat, x, w) w.axpy(1.0, y) def multTranspose(self, mat, x, y, inc=False): with self.ffn.dat.vec as v: x.copy(v) firedrake.restrict(self.ffn, self.cfn) for bc in self.cbcs: bc.zero(self.cfn) with self.cfn.dat.vec_ro as v: if inc: y.axpy(1.0, v) else: v.copy(y) def multTransposeAdd(self, mat, x, y, w): if y.handle == w.handle: self.multTranspose(mat, x, w, inc=True) else: self.multTranspose(mat, x, w) w.axpy(1.0, y) ctx = Interpolation(cfn, ffn, cbcs, fbcs) mat = PETSc.Mat().create() mat.setSizes(((nrow, None), (ncol, None))) mat.setType(mat.Type.PYTHON) mat.setPythonContext(ctx) mat.setUp() return mat, None