def reconstruct(self, field=None, V=None, subu=None, u=None, row_field=None, col_field=None, action_x=None, use_split=False): subu = subu or self.u row_field = row_field or field col_field = col_field or field # define W and form if field is None: # Returns self W = self._function_space form = self.f else: assert V is not None, "`V` can not be `None` when `field` is not `None`" W = self.as_subspace(field, V, use_split) if W is None: return rank = len(self.f.arguments()) splitter = ExtractSubBlock() if rank == 1: form = splitter.split(self.f, argument_indices=(row_field, )) elif rank == 2: form = splitter.split(self.f, argument_indices=(row_field, col_field)) if u is not None: form = replace(form, {self.u: u}) if action_x is not None: assert len(form.arguments()) == 2, "rank of self.f must be 2 when using action_x parameter" form = ufl_expr.action(form, action_x) ebc = EquationBCSplit(form, subu, self.sub_domain, method=self.method, V=W) for bc in self.bcs: if isinstance(bc, DirichletBC): ebc.add(bc.reconstruct(V=W, g=bc.function_arg, sub_domain=bc.sub_domain, method=bc.method, use_split=use_split)) elif isinstance(bc, EquationBCSplit): bc_temp = bc.reconstruct(field=field, V=V, subu=subu, u=u, row_field=row_field, col_field=col_field, action_x=action_x, use_split=use_split) # Due to the "if index", bc_temp can be None if bc_temp is not None: ebc.add(bc_temp) return ebc
def reconstruct(self, field=None, V=None, subu=None, u=None, row_field=None, col_field=None, action_x=None, use_split=False): subu = subu or self.u row_field = row_field or field col_field = col_field or field # define W and form if field is None: # Returns self W = self._function_space form = self.f else: assert V is not None, "`V` can not be `None` when `field` is not `None`" W = self.as_subspace(field, V, use_split) if W is None: return rank = len(self.f.arguments()) splitter = ExtractSubBlock() if rank == 1: form = splitter.split(self.f, argument_indices=(row_field, )) elif rank == 2: form = splitter.split(self.f, argument_indices=(row_field, col_field)) if u is not None: form = replace(form, {self.u: u}) if action_x is not None: assert len(form.arguments()) == 2, "rank of self.f must be 2 when using action_x parameter" form = ufl_expr.action(form, action_x) ebc = EquationBCSplit(form, subu, self.sub_domain, method=self.method, V=W) for bc in self.bcs: if isinstance(bc, DirichletBC): ebc.add(bc.reconstruct(V=W, g=bc.function_arg, sub_domain=bc.sub_domain, method=bc.method, use_split=use_split)) elif isinstance(bc, EquationBCSplit): bc_temp = bc.reconstruct(field=field, V=V, subu=subu, u=u, row_field=row_field, col_field=col_field, action_x=action_x, use_split=use_split) # Due to the "if index", bc_temp can be None if bc_temp is not None: ebc.add(bc_temp) return ebc
def split(self, fields): from ufl import as_vector, replace from firedrake import NonlinearVariationalProblem as NLVP, FunctionSpace splits = self._splits.get(tuple(fields)) if splits is not None: return splits splits = [] problem = self._problem splitter = ExtractSubBlock() for field in fields: try: if len(field) > 1: raise NotImplementedError("Can't split into subblock") except TypeError: # Just a single field, we can handle that pass F = splitter.split(problem.F, argument_indices=(field, )) J = splitter.split(problem.J, argument_indices=(field, field)) us = problem.u.split() subu = us[field] vec = [] for i, u in enumerate(us): for idx in numpy.ndindex(u.ufl_shape): vec.append(u[idx]) u = as_vector(vec) F = replace(F, {problem.u: u}) J = replace(J, {problem.u: u}) if problem.Jp is not None: Jp = splitter.split(problem.Jp, argument_indices=(field, field)) Jp = replace(Jp, {problem.u: u}) else: Jp = None bcs = [] for bc in problem.bcs: if bc.function_space().index == field: V = FunctionSpace(subu.ufl_domain(), subu.ufl_element()) bcs.append( type(bc)(V, bc.function_arg, bc.sub_domain, method=bc.method)) new_problem = NLVP( F, subu, bcs=bcs, J=J, Jp=None, form_compiler_parameters=problem.form_compiler_parameters) new_problem._constant_jacobian = problem._constant_jacobian splits.append( type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type, appctx=self.appctx)) return self._splits.setdefault(tuple(fields), splits)
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 split(self, fields): from ufl import as_vector, replace from firedrake import NonlinearVariationalProblem as NLVP, FunctionSpace splits = self._splits.get(tuple(fields)) if splits is not None: return splits splits = [] problem = self._problem splitter = ExtractSubBlock() for field in fields: try: if len(field) > 1: raise NotImplementedError("Can't split into subblock") except TypeError: # Just a single field, we can handle that pass F = splitter.split(problem.F, argument_indices=(field, )) J = splitter.split(problem.J, argument_indices=(field, field)) us = problem.u.split() subu = us[field] vec = [] for i, u in enumerate(us): for idx in numpy.ndindex(u.ufl_shape): vec.append(u[idx]) u = as_vector(vec) F = replace(F, {problem.u: u}) J = replace(J, {problem.u: u}) if problem.Jp is not None: Jp = splitter.split(problem.Jp, argument_indices=(field, field)) Jp = replace(Jp, {problem.u: u}) else: Jp = None bcs = [] for bc in problem.bcs: if bc.function_space().index == field: V = FunctionSpace(subu.ufl_domain(), subu.ufl_element()) bcs.append(type(bc)(V, bc.function_arg, bc.sub_domain, method=bc.method)) new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=None, form_compiler_parameters=problem.form_compiler_parameters) new_problem._constant_jacobian = problem._constant_jacobian splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type, appctx=self.appctx)) return self._splits.setdefault(tuple(fields), splits)
def test_split_coordinate_derivative(): mesh = UnitSquareMesh(1, 1) V = FunctionSpace(mesh, "P", 1) Q = FunctionSpace(mesh, "DP", 0) W = V * Q v = TestFunction(W) w = Function(W) x = SpatialCoordinate(mesh) J = derivative(inner(v, w) * dx, x) splitter = ExtractSubBlock() J00 = splitter.split(J, (0, 0)) expect = derivative(inner(as_vector([TestFunction(V), 0]), w) * dx, x) assert J00.signature() == expect.signature()
def test_split_coefficient_not_argument(): mesh = UnitSquareMesh(1, 1) V = FunctionSpace(mesh, "P", 1) Q = FunctionSpace(mesh, "DP", 0) W = V * Q w = Function(W) wr = Function(W) J = derivative(derivative(inner(grad(w), grad(w)) * dx, w), w, wr) splitter = ExtractSubBlock() J00 = splitter.split(J, (0, 0)) expect = derivative( derivative( inner(grad(w), grad(w)) * dx, w, as_vector([TestFunction(V), 0])), w, wr) assert J00.signature() == expect.signature()
def test_blocks(zero_rank_tensor, mixed_matrix, mixed_vector): S = zero_rank_tensor M = mixed_matrix F = mixed_vector a = M.form L = F.form splitter = ExtractSubBlock() M00 = M.block((0, 0)) M11 = M.block((1, 1)) M22 = M.block((2, 2)) M0101 = M.block(((0, 1), (0, 1))) M012 = M.block(((0, 1), (2, ))) M201 = M.block((((2, ), (0, 1)))) F0 = F.block((0, )) F1 = F.block((1, )) F2 = F.block((2, )) F01 = F.block(((0, 1), )) F12 = F.block(((1, 2), )) # Test index checking with pytest.raises(ValueError): S.block((0, )) with pytest.raises(ValueError): F.block((0, 1)) with pytest.raises(ValueError): M.block(((0, 1, 2, 3), 0)) with pytest.raises(ValueError): M.block((3, 3)) with pytest.raises(ValueError): F.block((3, )) # Check Tensor is (not) mixed where appropriate assert not M00.is_mixed assert not M11.is_mixed assert not M22.is_mixed assert not F0.is_mixed assert not F1.is_mixed assert not F2.is_mixed assert M0101.is_mixed assert M012.is_mixed assert M201.is_mixed assert F01.is_mixed assert F12.is_mixed # Taking blocks of non-mixed block (or scalars) should induce a no-op assert S.block(()) == S assert M00.block((0, 0)) == M00 assert M11.block((0, 0)) == M11 assert M22.block((0, 0)) == M22 assert F0.block((0, )) == F0 assert F1.block((0, )) == F1 assert F2.block((0, )) == F2 # Test arguments assert M00.arguments() == splitter.split(a, (0, 0)).arguments() assert M11.arguments() == splitter.split(a, (1, 1)).arguments() assert M22.arguments() == splitter.split(a, (2, 2)).arguments() assert F0.arguments() == splitter.split(L, (0, )).arguments() assert F1.arguments() == splitter.split(L, (1, )).arguments() assert F2.arguments() == splitter.split(L, (2, )).arguments() assert M0101.arguments() == splitter.split(a, ((0, 1), (0, 1))).arguments() assert M012.arguments() == splitter.split(a, ((0, 1), (2, ))).arguments() assert M201.arguments() == splitter.split(a, ((2, ), (0, 1))).arguments() assert F01.arguments() == splitter.split(L, ((0, 1), )).arguments() assert F12.arguments() == splitter.split(L, ((1, 2), )).arguments()
def split(self, fields): from firedrake import replace, as_vector, split from firedrake import NonlinearVariationalProblem as NLVP fields = tuple(tuple(f) for f in fields) splits = self._splits.get(tuple(fields)) if splits is not None: return splits splits = [] problem = self._problem splitter = ExtractSubBlock() for field in fields: F = splitter.split(problem.F, argument_indices=(field, )) J = splitter.split(problem.J, argument_indices=(field, field)) us = problem.u.split() V = F.arguments()[0].function_space() # Exposition: # We are going to make a new solution Function on the sub # mixed space defined by the relevant fields. # But the form may refer to the rest of the solution # anyway. # So we pull it apart and will make a new function on the # subspace that shares data. pieces = [us[i].dat for i in field] if len(pieces) == 1: val, = pieces subu = function.Function(V, val=val) subsplit = (subu, ) else: val = op2.MixedDat(pieces) subu = function.Function(V, val=val) # Split it apart to shove in the form. subsplit = split(subu) # Permutation from field indexing to indexing of pieces field_renumbering = dict([f, i] for i, f in enumerate(field)) vec = [] for i, u in enumerate(us): if i in field: # If this is a field we're keeping, get it from # the new function. Otherwise just point to the # old data. u = subsplit[field_renumbering[i]] if u.ufl_shape == (): vec.append(u) else: for idx in numpy.ndindex(u.ufl_shape): vec.append(u[idx]) # So now we have a new representation for the solution # vector in the old problem. For the fields we're going # to solve for, it points to a new Function (which wraps # the original pieces). For the rest, it points to the # pieces from the original Function. # IOW, we've reinterpreted our original mixed solution # function as being made up of some spaces we're still # solving for, and some spaces that have just become # coefficients in the new form. u = as_vector(vec) F = replace(F, {problem.u: u}) J = replace(J, {problem.u: u}) if problem.Jp is not None: Jp = splitter.split(problem.Jp, argument_indices=(field, field)) Jp = replace(Jp, {problem.u: u}) else: Jp = None bcs = [] for bc in problem.bcs: Vbc = bc.function_space() if Vbc.parent is not None and isinstance(Vbc.parent.ufl_element(), VectorElement): index = Vbc.parent.index else: index = Vbc.index cmpt = Vbc.component # TODO: need to test this logic if index in field: if len(field) == 1: W = V else: W = V.sub(field_renumbering[index]) if cmpt is not None: W = W.sub(cmpt) bcs.append(type(bc)(W, bc.function_arg, bc.sub_domain, method=bc.method)) new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=Jp, form_compiler_parameters=problem.form_compiler_parameters) new_problem._constant_jacobian = problem._constant_jacobian splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type, appctx=self.appctx)) return self._splits.setdefault(tuple(fields), splits)
def split(self, fields): from firedrake import replace, as_vector, split from firedrake_ts.ts_solver import DAEProblem from firedrake.bcs import DirichletBC, EquationBC fields = tuple(tuple(f) for f in fields) splits = self._splits.get(tuple(fields)) if splits is not None: return splits splits = [] problem = self._problem splitter = ExtractSubBlock() for field in fields: F = splitter.split(problem.F, argument_indices=(field, )) J = splitter.split(problem.J, argument_indices=(field, field)) us = problem.u.split() V = F.arguments()[0].function_space() # Exposition: # We are going to make a new solution Function on the sub # mixed space defined by the relevant fields. # But the form may refer to the rest of the solution # anyway. # So we pull it apart and will make a new function on the # subspace that shares data. pieces = [us[i].dat for i in field] if len(pieces) == 1: (val, ) = pieces subu = function.Function(V, val=val) subsplit = (subu, ) else: val = op2.MixedDat(pieces) subu = function.Function(V, val=val) # Split it apart to shove in the form. subsplit = split(subu) # Permutation from field indexing to indexing of pieces field_renumbering = dict([f, i] for i, f in enumerate(field)) vec = [] for i, u in enumerate(us): if i in field: # If this is a field we're keeping, get it from # the new function. Otherwise just point to the # old data. u = subsplit[field_renumbering[i]] if u.ufl_shape == (): vec.append(u) else: for idx in numpy.ndindex(u.ufl_shape): vec.append(u[idx]) # So now we have a new representation for the solution # vector in the old problem. For the fields we're going # to solve for, it points to a new Function (which wraps # the original pieces). For the rest, it points to the # pieces from the original Function. # IOW, we've reinterpreted our original mixed solution # function as being made up of some spaces we're still # solving for, and some spaces that have just become # coefficients in the new form. u = as_vector(vec) F = replace(F, {problem.u: u}) J = replace(J, {problem.u: u}) if problem.Jp is not None: Jp = splitter.split(problem.Jp, argument_indices=(field, field)) Jp = replace(Jp, {problem.u: u}) else: Jp = None bcs = [] for bc in problem.bcs: if isinstance(bc, DirichletBC): bc_temp = bc.reconstruct( field=field, V=V, g=bc.function_arg, sub_domain=bc.sub_domain, method=bc.method, ) elif isinstance(bc, EquationBC): bc_temp = bc.reconstruct(field, V, subu, u) if bc_temp is not None: bcs.append(bc_temp) new_problem = DAEProblem( F, subu, problem.udot, problem.tspan, bcs=bcs, J=J, Jp=Jp, form_compiler_parameters=problem.form_compiler_parameters, ) new_problem._constant_jacobian = problem._constant_jacobian splits.append( type(self)( new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type, appctx=self.appctx, transfer_manager=self.transfer_manager, )) return self._splits.setdefault(tuple(fields), splits)
def test_blocks(zero_rank_tensor, mixed_matrix, mixed_vector): S = zero_rank_tensor M = mixed_matrix F = mixed_vector a = M.form L = F.form splitter = ExtractSubBlock() _M = M.blocks M00 = _M[0, 0] M11 = _M[1, 1] M22 = _M[2, 2] M0101 = _M[:2, :2] M012 = _M[:2, 2] M201 = _M[2, :2] _F = F.blocks F0 = _F[0] F1 = _F[1] F2 = _F[2] F01 = _F[:2] F12 = _F[1:3] # Test index checking with pytest.raises(ValueError): S.blocks[0] with pytest.raises(ValueError): _F[0, 1] with pytest.raises(ValueError): _M[0:5, 2] with pytest.raises(ValueError): _M[3, 3] with pytest.raises(ValueError): _F[3] # Check Tensor is (not) mixed where appropriate assert not M00.is_mixed assert not M11.is_mixed assert not M22.is_mixed assert not F0.is_mixed assert not F1.is_mixed assert not F2.is_mixed assert M0101.is_mixed assert M012.is_mixed assert M201.is_mixed assert F01.is_mixed assert F12.is_mixed # Taking blocks of non-mixed block (or scalars) should induce a no-op assert S.blocks[None] == S # This is silly, but it's technically a no-op assert M00.blocks[0, 0] == M00 assert M11.blocks[0, 0] == M11 assert M22.blocks[0, 0] == M22 assert F0.blocks[0] == F0 assert F1.blocks[0] == F1 assert F2.blocks[0] == F2 # Test arguments assert M00.arguments() == splitter.split(a, (0, 0)).arguments() assert M11.arguments() == splitter.split(a, (1, 1)).arguments() assert M22.arguments() == splitter.split(a, (2, 2)).arguments() assert F0.arguments() == splitter.split(L, (0, )).arguments() assert F1.arguments() == splitter.split(L, (1, )).arguments() assert F2.arguments() == splitter.split(L, (2, )).arguments() assert M0101.arguments() == splitter.split(a, ((0, 1), (0, 1))).arguments() assert M012.arguments() == splitter.split(a, ((0, 1), (2, ))).arguments() assert M201.arguments() == splitter.split(a, ((2, ), (0, 1))).arguments() assert F01.arguments() == splitter.split(L, ((0, 1), )).arguments() assert F12.arguments() == splitter.split(L, ((1, 2), )).arguments()