def compute_shapes(form: ufl.Form): """Computes the shapes of the cell tensor, coefficient & coordinate arrays for a form""" # Cell tensor elements = tuple(arg.ufl_element() for arg in form.arguments()) fiat_elements = (create_element(element) for element in elements) element_dims = tuple(fe.space_dimension() for fe in fiat_elements) A_shape = element_dims # Coefficient arrays ws_shapes = [] for coefficient in form.coefficients(): w_element = coefficient.ufl_element() w_dim = create_element(w_element).space_dimension() ws_shapes.append((w_dim,)) if len(ws_shapes) > 0: max_w_dim = sorted(ws_shapes, key=lambda w_dim: np.prod(w_dim))[-1] w_full_shape = (len(ws_shapes), ) + max_w_dim else: w_full_shape = (0,) # Coordinate dofs array num_vertices_per_cell = form.ufl_cell().num_vertices() gdim = form.ufl_cell().geometric_dimension() coords_shape = (num_vertices_per_cell, gdim) return A_shape, w_full_shape, coords_shape
def __init__(self, form: ufl.Form, form_compiler_parameters: dict = {}, jit_parameters: dict = {}): """Create dolfinx Form Parameters ---------- form Pure UFL form form_compiler_parameters See :py:func:`ffcx_jit <dolfinx.jit.ffcx_jit>` jit_parameters See :py:func:`ffcx_jit <dolfinx.jit.ffcx_jit>` Note ---- This wrapper for UFL form is responsible for the actual FFCX compilation and attaching coefficients and domains specific data to the underlying C++ Form. """ # Extract subdomain data from UFL form sd = form.subdomain_data() self._subdomains, = list(sd.values()) # Assuming single domain domain, = list(sd.keys()) # Assuming single domain mesh = domain.ufl_cargo() if mesh is None: raise RuntimeError("Expecting to find a Mesh in the form.") # Compile UFL form with JIT ufc_form = jit.ffcx_jit( mesh.mpi_comm(), form, form_compiler_parameters=form_compiler_parameters, jit_parameters=jit_parameters) # For every argument in form extract its function space function_spaces = [ func.ufl_function_space()._cpp_object for func in form.arguments() ] # Prepare coefficients data. For every coefficient in form take # its C++ object. original_coefficients = form.coefficients() coeffs = [original_coefficients[ufc_form.original_coefficient_position( i)]._cpp_object for i in range(ufc_form.num_coefficients)] # Create dictionary of of subdomain markers (possible None for # some dimensions subdomains = {cpp.fem.IntegralType.cell: self._subdomains.get("cell"), cpp.fem.IntegralType.exterior_facet: self._subdomains.get("exterior_facet"), cpp.fem.IntegralType.interior_facet: self._subdomains.get("interior_facet"), cpp.fem.IntegralType.vertex: self._subdomains.get("vertex")} # Prepare dolfinx.cpp.fem.Form and hold it as a member ffi = cffi.FFI() self._cpp_object = cpp.fem.create_form(ffi.cast("uintptr_t", ufc_form), function_spaces, coeffs, [c._cpp_object for c in form.constants()], subdomains, mesh)
def compile_terminal_form(tensor, prefix=None, tsfc_parameters=None, coffee=True): """Compiles the TSFC form associated with a Slate :class:`Tensor` object. This function will return a :class:`ContextKernel` which stores information about the original tensor, integral types and the corresponding TSFC kernels. :arg tensor: A Slate `Tensor`. :arg prefix: An optional `string` indicating the prefix for the subkernel. :arg tsfc_parameters: An optional `dict` of parameters to provide TSFC. Returns: A `ContextKernel` containing all relevant information. """ assert isinstance( tensor, Tensor), ("Only terminal tensors have forms associated with them!") # Sets a default name for the subkernel prefix. # NOTE: the builder will choose a prefix independent of # the tensor name for code idempotency reasons, but is not # strictly required. prefix = prefix or "subkernel%s_" % tensor.__str__() mapper = RemoveRestrictions() integrals = map(partial(map_integrand_dags, mapper), tensor.form.integrals()) transformed_integrals = transform_integrals(integrals) cxt_kernels = [] for orig_it_type, integrals in transformed_integrals.items(): subkernel_prefix = prefix + "%s_to_" % orig_it_type form = Form(integrals) kernels = tsfc_compile(form, subkernel_prefix, parameters=tsfc_parameters, coffee=coffee, split=False) if kernels: cxt_k = ContextKernel(tensor=tensor, coefficients=form.coefficients(), original_integral_type=orig_it_type, tsfc_kernels=kernels) cxt_kernels.append(cxt_k) cxt_kernels = tuple(cxt_kernels) return cxt_kernels
def prepare_tsfc_kernels(temps, tsfc_parameters=None): """This function generates a mapping of the form: ``kernel_exprs = {terminal_node: kernels}`` where `terminal_node` objects are :class:`slate.Tensor` nodes and `kernels` is an iterable of `namedtuple` objects, `SplitKernel`, provided by TSFC. This mapping is used in :class:`SlateKernelBuilder` to provide direct access to all `SplitKernel` objects associated with a `slate.Tensor` node. :arg temps: a mapping of the form ``{terminal_node: symbol_name}`` (see :meth:`prepare_temporaries`). :arg tsfc_parameters: an optional `dict` of parameters to pass onto TSFC. """ kernel_exprs = {} for expr in temps.keys(): integrals = expr.form.integrals() mapper = RemoveRestrictions() integrals = map(partial(map_integrand_dags, mapper), integrals) prefix = "subkernel%d_" % len(kernel_exprs) # Now we split integrals by type: interior_facet and all other cases # First, the interior_facet case: interior_facet_intergrals = filter( lambda x: x.integral_type() == "interior_facet", integrals) # Now we reconstruct all interior_facet integrals to be of type: exterior_facet # This is because locally over each cell, SLATE views them as being "exterior" # with respect to the cell. interior_facet_intergrals = [ it.reconstruct(integral_type="exterior_facet") for it in interior_facet_intergrals ] # Now for the rest: other_integrals = filter( lambda x: x.integral_type() != "interior_facet", integrals) forms = (Form(interior_facet_intergrals), Form(other_integrals)) compiled_forms = [] for form in forms: compiled_forms.extend( compile_form(form, prefix, parameters=tsfc_parameters)) kernel_exprs[expr] = tuple(compiled_forms) return kernel_exprs
def vjp_assemble_impl( g: np.array, fenics_output_form: ufl.Form, fenics_inputs: List[FenicsVariable], ) -> Tuple[np.array]: """Computes the gradients of the output with respect to the inputs.""" # Compute derivative form for the output with respect to each input fenics_grads_forms = [] for fenics_input in fenics_inputs: # Need to construct direction (test function) first if isinstance(fenics_input, fenics.Function): V = fenics_input.function_space() elif isinstance(fenics_input, fenics.Constant): mesh = fenics_output_form.ufl_domain().ufl_cargo() V = fenics.FunctionSpace(mesh, "Real", 0) else: raise NotImplementedError dv = fenics.TestFunction(V) fenics_grad_form = fenics.derivative(fenics_output_form, fenics_input, dv) fenics_grads_forms.append(fenics_grad_form) # Assemble the derivative forms fenics_grads = [fenics.assemble(form) for form in fenics_grads_forms] # Convert FEniCS gradients to jax array representation jax_grads = (None if fg is None else np.asarray(g * fenics_to_numpy(fg)) for fg in fenics_grads) jax_grad_tuple = tuple(jax_grads) return jax_grad_tuple
def adjoint(form: ufl.Form, reordered_arguments=None) -> ufl.Form: """Compute adjoint of a bilinear form by changing the ordering (count) of the test and trial functions. The functions wraps ``ufl.adjoint``, and by default UFL will create new ``Argument`` s. To specify the ``Argument`` s rather than creating new ones, pass a tuple of ``Argument`` s as ``reordered_arguments``. See the documentation for ``ufl.adjoint`` for more details. """ if reordered_arguments is not None: return ufl.adjoint(form, reordered_arguments=reordered_arguments) # Extract form arguments arguments = form.arguments() if len(arguments) != 2: raise RuntimeError( "Cannot compute adjoint of form, form is not bilinear") if any(arg.part() is not None for arg in arguments): raise RuntimeError( "Cannot compute adjoint of form, parts not supported") # Create new Arguments in the same spaces (NB: Order does not matter # anymore here because number is absolute) v1 = function.Argument(arguments[1].function_space, arguments[0].number(), arguments[0].part()) v0 = function.Argument(arguments[0].function_space, arguments[1].number(), arguments[1].part()) # Return form with swapped arguments as new arguments return ufl.adjoint(form, reordered_arguments=(v1, v0))
def coarsen_form(form, Nf, Nc, replace_d): # Coarsen a form, by replacing the solution, test and trial functions, and # reconstructing each integral with a coarsened quadrature degree. # If form is not a Form, then return form. return Form([ f.reconstruct( metadata=coarsen_quadrature(f.metadata(), Nf, Nc)) for f in replace(form, replace_d).integrals() ]) if isinstance(form, Form) else form
def sum_integrands(form): """Produce a form with the integrands on the same measure summed.""" integrals = defaultdict(list) for integral in form.integrals(): md = integral.metadata() mdkey = tuple((k, md[k]) for k in sorted(md.keys())) integrals[(integral.integral_type(), integral.ufl_domain(), integral.subdomain_data(), integral.subdomain_id(), mdkey)].append(integral) return Form([ it[0].reconstruct(reduce(add, [i.integrand() for i in it])) for it in integrals.values() ])
def reconstruct_form_from_sub_integral_data(sub_integral_data, domain_data=None): domain_data = domain_data or {} integrals = [] # Iterate over domain types in order for dt in Measure._domain_types_tuple: dd = domain_data.get(dt) # Get integrals list for this domain type if any metaintegrands = sub_integral_data.get(dt) if metaintegrands is not None: for k in sorted(metaintegrands.keys()): for integrand, compiler_data in metaintegrands[k]: integrals.append( Integral(integrand, dt, k, compiler_data, dd)) return Form(integrals)
def prune(self, form, check_zeros=True): # Get the parts of the form with the correct arity if self._index_u is None: form = compute_form_rhs(form) else: form = compute_form_lhs(form) integrals = [] for integral in form.integrals(): # Prune integrals that do not contain Arguments with # the chosen coupled function space indices pruned = self.visit(integral.integrand()) if not check_zeros or not is_zero_ufl_expression(pruned): integrals.append(integral.reconstruct(pruned)) return Form(integrals)
def sum_integrands(form): """Produce a form with the integrands on the same measure summed.""" integrals = defaultdict(list) for integral in form.integrals(): md = integral.metadata() mdkey = tuple((k, md[k]) for k in sorted(md.keys())) sd = integral.subdomain_data() # subdomain data is a weakref sd = sd() assert sd is not None, "Lost reference to subdomain data, argh!" integrals[(integral.integral_type(), integral.domain(), integral.subdomain_id(), sd, mdkey)].append(integral) return Form([ it[0].reconstruct(reduce(add, [i.integrand() for i in it])) for it in integrals.values() ])
def expand_sum_product(form): form = remove_complex_nodes(form) # TODO support forms in the complex field. This is currently needed otherwise conj((a+b)*c) does not get expanded. # Patch Expr.__mul__ and Expr.__rmul__ patch_expr_mul() # Call sympy replacer expanded_form = map_integrand_dags(SympyExpander(), form) # Split sums expanded_split_form_integrals = list() for integral in expanded_form.integrals(): expanded_split_form_integrands = list() split_sum(integral.integrand(), expanded_split_form_integrands) expanded_split_form_integrals.extend([integral.reconstruct(integrand=integrand) for integrand in expanded_split_form_integrands]) expanded_split_form = Form(expanded_split_form_integrals) # Undo patch to Expr.__mul__ and Expr.__rmul__ unpatch_expr_mul() # Return return expanded_split_form
def expand_sum_product(form): # Patch Expr.__mul__ and Expr.__rmul__ patch_expr_mul() # Call sympy replacer expanded_form = map_integrand_dags(SympyExpander(), form) # Split sums expanded_split_form_integrals = list() for integral in expanded_form.integrals(): expanded_split_form_integrands = list() split_sum(integral.integrand(), expanded_split_form_integrands) expanded_split_form_integrals.extend([ integral.reconstruct(integrand=integrand) for integrand in expanded_split_form_integrands ]) expanded_split_form = Form(expanded_split_form_integrals) # Undo patch to Expr.__mul__ and Expr.__rmul__ unpatch_expr_mul() # Return return expanded_split_form
def __init__(self, a: ufl.Form, L: ufl.Form, bcs: typing.List[fem.DirichletBC] = [], u: fem.Function = None, petsc_options={}, form_compiler_parameters={}, jit_parameters={}): """Initialize solver for a linear variational problem. Parameters ---------- a A bilinear UFL form, the left hand side of the variational problem. L A linear UFL form, the right hand side of the variational problem. bcs A list of Dirichlet boundary conditions. u The solution function. It will be created if not provided. petsc_options Parameters that is passed to the linear algebra backend PETSc. For available choices for the 'petsc_options' kwarg, see the `PETSc-documentation <https://www.mcs.anl.gov/petsc/documentation/index.html>`. form_compiler_parameters Parameters used in FFCx compilation of this form. Run `ffcx --help` at the commandline to see all available options. Takes priority over all other parameter values, except for `scalar_type` which is determined by DOLFINx. jit_parameters Parameters used in CFFI JIT compilation of C code generated by FFCx. See `python/dolfinx/jit.py` for all available parameters. Takes priority over all other parameter values. .. code-block:: python problem = LinearProblem(a, L, [bc0, bc1], petsc_options={"ksp_type": "preonly", "pc_type": "lu"}) """ self._a = fem.Form(a, form_compiler_parameters=form_compiler_parameters, jit_parameters=jit_parameters) self._A = fem.create_matrix(self._a) self._L = fem.Form(L, form_compiler_parameters=form_compiler_parameters, jit_parameters=jit_parameters) self._b = fem.create_vector(self._L) if u is None: # Extract function space from TrialFunction (which is at the # end of the argument list as it is numbered as 1, while the # Test function is numbered as 0) self.u = fem.Function(a.arguments()[-1].ufl_function_space()) else: self.u = u self.bcs = bcs self._solver = PETSc.KSP().create( self.u.function_space.mesh.mpi_comm()) self._solver.setOperators(self._A) # Give PETSc solver options a unique prefix solver_prefix = "dolfinx_solve_{}".format(id(self)) self._solver.setOptionsPrefix(solver_prefix) # Set PETSc options opts = PETSc.Options() opts.prefixPush(solver_prefix) for k, v in petsc_options.items(): opts[k] = v opts.prefixPop() self._solver.setFromOptions()
def __init__(self, form: ufl.Form, form_compiler_parameters: dict = {}, jit_parameters: dict = {}): """Create dolfinx Form Parameters ---------- form Pure UFL form form_compiler_parameters Parameters used in FFCX compilation of this form. Run `ffcx --help` in the commandline to see all available options. jit_parameters Parameters controlling JIT compilation of C code. Note ---- This wrapper for UFL form is responsible for the actual FFCX compilation and attaching coefficients and domains specific data to the underlying C++ Form. """ # Extract subdomain data from UFL form sd = form.subdomain_data() self._subdomains, = list(sd.values()) # Assuming single domain domain, = list(sd.keys()) # Assuming single domain mesh = domain.ufl_cargo() # Compile UFL form with JIT ufc_form = jit.ffcx_jit( form, form_compiler_parameters=form_compiler_parameters, jit_parameters=jit_parameters, mpi_comm=mesh.mpi_comm()) # For every argument in form extract its function space function_spaces = [ func.ufl_function_space()._cpp_object for func in form.arguments() ] # Prepare dolfinx.Form and hold it as a member ffi = cffi.FFI() self._cpp_object = cpp.fem.create_form(ffi.cast("uintptr_t", ufc_form), function_spaces) # Need to fill the form with coefficients data # For every coefficient in form take its C++ object original_coefficients = form.coefficients() for i in range(self._cpp_object.num_coefficients()): j = self._cpp_object.original_coefficient_position(i) self._cpp_object.set_coefficient( i, original_coefficients[j]._cpp_object) # Constants are set based on their position in original form original_constants = [c._cpp_object for c in form.constants()] self._cpp_object.set_constants(original_constants) if mesh is None: raise RuntimeError("Expecting to find a Mesh in the form.") # Attach mesh (because function spaces and coefficients may be # empty lists) if not function_spaces: self._cpp_object.set_mesh(mesh) # Attach subdomains to C++ Form if we have them subdomains = self._subdomains.get("cell") if subdomains: self._cpp_object.set_cell_domains(subdomains) subdomains = self._subdomains.get("exterior_facet") if subdomains: self._cpp_object.set_exterior_facet_domains(subdomains) subdomains = self._subdomains.get("interior_facet") if subdomains: self._cpp_object.set_interior_facet_domains(subdomains) subdomains = self._subdomains.get("vertex") if subdomains: self._cpp_object.set_vertex_domains(subdomains)
def __init__(self, a: ufl.Form, L: ufl.Form, bcs: typing.List[DirichletBCMetaClass] = [], u: _Function = None, petsc_options={}, form_compiler_params={}, jit_params={}): """Initialize solver for a linear variational problem. Args: a: A bilinear UFL form, the left hand side of the variational problem. L: A linear UFL form, the right hand side of the variational problem. bcs: A list of Dirichlet boundary conditions. u: The solution function. It will be created if not provided. petsc_options: Parameters that is passed to the linear algebra backend PETSc. For available choices for the 'petsc_options' kwarg, see the `PETSc documentation <https://petsc4py.readthedocs.io/en/stable/manual/ksp/>`_. form_compiler_params: Parameters used in FFCx compilation of this form. Run ``ffcx --help`` at the commandline to see all available options. jit_params: Parameters used in CFFI JIT compilation of C code generated by FFCx. See `python/dolfinx/jit.py` for all available parameters. Takes priority over all other parameter values. Example:: problem = LinearProblem(a, L, [bc0, bc1], petsc_options={"ksp_type": "preonly", "pc_type": "lu"}) """ self._a = _create_form(a, form_compiler_params=form_compiler_params, jit_params=jit_params) self._A = create_matrix(self._a) self._L = _create_form(L, form_compiler_params=form_compiler_params, jit_params=jit_params) self._b = create_vector(self._L) if u is None: # Extract function space from TrialFunction (which is at the # end of the argument list as it is numbered as 1, while the # Test function is numbered as 0) self.u = _Function(a.arguments()[-1].ufl_function_space()) else: self.u = u self._x = _cpp.la.petsc.create_vector_wrap(self.u.x) self.bcs = bcs self._solver = PETSc.KSP().create(self.u.function_space.mesh.comm) self._solver.setOperators(self._A) # Give PETSc solver options a unique prefix problem_prefix = "dolfinx_solve_{}".format(id(self)) self._solver.setOptionsPrefix(problem_prefix) # Set PETSc options opts = PETSc.Options() opts.prefixPush(problem_prefix) for k, v in petsc_options.items(): opts[k] = v opts.prefixPop() self._solver.setFromOptions() # Set matrix and vector PETSc options self._A.setOptionsPrefix(problem_prefix) self._A.setFromOptions() self._b.setOptionsPrefix(problem_prefix) self._b.setFromOptions()
def __init__(self, form: ufl.Form, form_compiler_parameters: dict = None): """Create dolfin Form Parameters ---------- form Pure UFL form form_compiler_parameters Parameters used in JIT FFC compilation of this form Note ---- This wrapper for UFL form is responsible for the actual FFC compilation and attaching coefficients and domains specific data to the underlying C++ Form. """ self.form_compiler_parameters = form_compiler_parameters # Extract subdomain data from UFL form sd = form.subdomain_data() self._subdomains, = list(sd.values()) # Assuming single domain domain, = list(sd.keys()) # Assuming single domain mesh = domain.ufl_cargo() # Compile UFL form with JIT ufc_form = jit.ffc_jit( form, form_compiler_parameters=self.form_compiler_parameters, mpi_comm=mesh.mpi_comm()) # Cast compiled library to pointer to ufc_form ffi = cffi.FFI() ufc_form = fem.dofmap.make_ufc_form(ffi.cast("uintptr_t", ufc_form)) # For every argument in form extract its function space function_spaces = [ func.function_space()._cpp_object for func in form.arguments() ] # Prepare dolfin.Form and hold it as a member self._cpp_object = cpp.fem.create_form(ufc_form, function_spaces) # Need to fill the form with coefficients data # For every coefficient in form take its CPP object original_coefficients = form.coefficients() for i in range(self._cpp_object.num_coefficients()): j = self._cpp_object.original_coefficient_position(i) self._cpp_object.set_coefficient( i, original_coefficients[j]._cpp_object) if mesh is None: raise RuntimeError("Expecting to find a Mesh in the form.") # Attach mesh (because function spaces and coefficients may be # empty lists) if not function_spaces: self._cpp_object.set_mesh(mesh) # Attach subdomains to C++ Form if we have them subdomains = self._subdomains.get("cell") if subdomains: self._cpp_object.set_cell_domains(subdomains) subdomains = self._subdomains.get("exterior_facet") if subdomains: self._cpp_object.set_exterior_facet_domains(subdomains) subdomains = self._subdomains.get("interior_facet") if subdomains: self._cpp_object.set_interior_facet_domains(subdomains) subdomains = self._subdomains.get("vertex") if subdomains: self._cpp_object.set_vertex_domains(subdomains)
def __init__(self, form: ufl.Form, form_compiler_parameters: dict = None): """Create dolfin Form Parameters ---------- form Pure UFL form form_compiler_parameters Parameters used in JIT FFC compilation of this form Note ---- This wrapper for UFL form is responsible for the actual FFC compilation and attaching coefficients and domains specific data to the underlying C++ Form. """ self.form_compiler_parameters = form_compiler_parameters # Add DOLFIN include paths (just the Boost path for special # math functions is really required) # FIXME: move getting include paths to elsewhere if self.form_compiler_parameters is None: self.form_compiler_parameters = { "external_include_dirs": jit.dolfin_pc["include_dirs"] } else: # FIXME: add paths if dict entry already exists self.form_compiler_parameters[ "external_include_dirs"] = jit.dolfin_pc["include_dirs"] # Extract subdomain data from UFL form sd = form.subdomain_data() self._subdomains, = list(sd.values()) # Assuming single domain domain, = list(sd.keys()) # Assuming single domain mesh = domain.ufl_cargo() # Compile UFL form with JIT ufc_form = jit.ffc_jit( form, form_compiler_parameters=self.form_compiler_parameters, mpi_comm=mesh.mpi_comm()) # Cast compiled library to pointer to ufc_form ufc_form = fem.dofmap.make_ufc_form(ufc_form[0]) # For every argument in form extract its function space function_spaces = [ func.function_space()._cpp_object for func in form.arguments() ] # Prepare dolfin.Form and hold it as a member self._cpp_object = cpp.fem.Form(ufc_form, function_spaces) # Need to fill the form with coefficients data # For every coefficient in form take its CPP object original_coefficients = form.coefficients() for i in range(self._cpp_object.num_coefficients()): j = self._cpp_object.original_coefficient_position(i) self._cpp_object.set_coefficient( j, original_coefficients[i]._cpp_object) if mesh is None: raise RuntimeError("Expecting to find a Mesh in the form.") # Attach mesh (because function spaces and coefficients may be # empty lists) if not function_spaces: self._cpp_object.set_mesh(mesh) # Attach subdomains to C++ Form if we have them subdomains = self._subdomains.get("cell") self._cpp_object.set_cell_domains(subdomains) subdomains = self._subdomains.get("exterior_facet") self._cpp_object.set_exterior_facet_domains(subdomains) subdomains = self._subdomains.get("interior_facet") self._cpp_object.set_interior_facet_domains(subdomains) subdomains = self._subdomains.get("vertex") self._cpp_object.set_vertex_domains(subdomains)