Spaces:
Paused
Paused
| import warp as wp | |
| from warp.fem.space import SpaceRestriction, FunctionSpace | |
| from warp.fem.types import Sample, get_node_index_in_element | |
| from warp.fem import utils, cache | |
| from .field import SpaceField | |
| class TestField(SpaceField): | |
| """Field defined over a space restriction that can be used as a test function. | |
| In order to reuse computations, it is possible to define the test field using a SpaceRestriction | |
| defined for a different value type than the test function value type, as long as the node topology is similar. | |
| """ | |
| def __init__(self, space_restriction: SpaceRestriction, space: FunctionSpace): | |
| if space_restriction.domain.dimension == space.dimension - 1: | |
| space = space.trace() | |
| if space_restriction.domain.dimension != space.dimension: | |
| raise ValueError("Incompatible space and domain dimensions") | |
| if space.topology != space_restriction.space_topology: | |
| raise ValueError("Incompatible space and space partition topologies") | |
| super().__init__(space, space_restriction.space_partition) | |
| self.space_restriction = space_restriction | |
| self.domain = self.space_restriction.domain | |
| self.EvalArg = self.space.SpaceArg | |
| self.ElementEvalArg = self._make_element_eval_arg() | |
| self.eval_degree = self._make_eval_degree() | |
| self.eval_inner = self._make_eval_inner() | |
| self.eval_grad_inner = self._make_eval_grad_inner() | |
| self.eval_div_inner = self._make_eval_div_inner() | |
| self.eval_outer = self._make_eval_outer() | |
| self.eval_grad_outer = self._make_eval_grad_outer() | |
| self.eval_div_outer = self._make_eval_div_outer() | |
| self.at_node = self._make_at_node() | |
| def name(self) -> str: | |
| return self.space.name + "Test" | |
| def eval_arg_value(self, device) -> wp.codegen.StructInstance: | |
| return self.space.space_arg_value(device) | |
| def _make_element_eval_arg(self): | |
| from warp.fem import cache | |
| class ElementEvalArg: | |
| elt_arg: self.domain.ElementArg | |
| eval_arg: self.EvalArg | |
| return ElementEvalArg | |
| def _make_eval_inner(self): | |
| def eval_test_inner(args: self.ElementEvalArg, s: Sample): | |
| weight = self.space.element_inner_weight( | |
| args.elt_arg, | |
| args.eval_arg, | |
| s.element_index, | |
| s.element_coords, | |
| get_node_index_in_element(s.test_dof), | |
| ) | |
| return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof) | |
| return eval_test_inner | |
| def _make_eval_grad_inner(self): | |
| if not self.gradient_valid(): | |
| return None | |
| def eval_nabla_test_inner(args: self.ElementEvalArg, s: Sample): | |
| nabla_weight = self.space.element_inner_weight_gradient( | |
| args.elt_arg, | |
| args.eval_arg, | |
| s.element_index, | |
| s.element_coords, | |
| get_node_index_in_element(s.test_dof), | |
| ) | |
| grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s) | |
| return utils.generalized_outer( | |
| self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof), | |
| utils.apply_right(nabla_weight, grad_transform), | |
| ) | |
| return eval_nabla_test_inner | |
| def _make_eval_div_inner(self): | |
| if not self.divergence_valid(): | |
| return None | |
| def eval_div_test_inner(args: self.ElementEvalArg, s: Sample): | |
| nabla_weight = self.space.element_inner_weight_gradient( | |
| args.elt_arg, | |
| args.eval_arg, | |
| s.element_index, | |
| s.element_coords, | |
| get_node_index_in_element(s.test_dof), | |
| ) | |
| grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s) | |
| return utils.generalized_inner( | |
| self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof), | |
| utils.apply_right(nabla_weight, grad_transform), | |
| ) | |
| return eval_div_test_inner | |
| def _make_eval_outer(self): | |
| def eval_test_outer(args: self.ElementEvalArg, s: Sample): | |
| weight = self.space.element_outer_weight( | |
| args.elt_arg, | |
| args.eval_arg, | |
| s.element_index, | |
| s.element_coords, | |
| get_node_index_in_element(s.test_dof), | |
| ) | |
| return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof) | |
| return eval_test_outer | |
| def _make_eval_grad_outer(self): | |
| if not self.gradient_valid(): | |
| return None | |
| def eval_nabla_test_outer(args: self.ElementEvalArg, s: Sample): | |
| nabla_weight = self.space.element_outer_weight_gradient( | |
| args.elt_arg, | |
| args.eval_arg, | |
| s.element_index, | |
| s.element_coords, | |
| get_node_index_in_element(s.test_dof), | |
| ) | |
| grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s) | |
| return utils.generalized_outer( | |
| self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof), | |
| utils.apply_right(nabla_weight, grad_transform), | |
| ) | |
| return eval_nabla_test_outer | |
| def _make_eval_div_outer(self): | |
| if not self.divergence_valid(): | |
| return None | |
| def eval_div_test_outer(args: self.ElementEvalArg, s: Sample): | |
| nabla_weight = self.space.element_outer_weight_gradient( | |
| args.elt_arg, | |
| args.eval_arg, | |
| s.element_index, | |
| s.element_coords, | |
| get_node_index_in_element(s.test_dof), | |
| ) | |
| grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s) | |
| return utils.generalized_inner( | |
| self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof), | |
| utils.apply_right(nabla_weight, grad_transform), | |
| ) | |
| return eval_div_test_outer | |
| def _make_at_node(self): | |
| def at_node(args: self.ElementEvalArg, s: Sample): | |
| node_coords = self.space.node_coords_in_element( | |
| args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(s.test_dof) | |
| ) | |
| return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof) | |
| return at_node | |