warp-lang 1.4.1__py3-none-manylinux2014_aarch64.whl → 1.5.0__py3-none-manylinux2014_aarch64.whl

warp/fem/field/test.py

@@ -1,180 +0,0 @@
-import warp as wp
-from warp.fem import cache, utils
-from warp.fem.space import FunctionSpace, SpaceRestriction
-from warp.fem.types import Sample, get_node_index_in_element
-
-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()
-
-    @property
-    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
-
-        @cache.dynamic_struct(suffix=self.name)
-        class ElementEvalArg:
-            elt_arg: self.domain.ElementArg
-            eval_arg: self.EvalArg
-
-        return ElementEvalArg
-
-    def _make_eval_inner(self):
-        @cache.dynamic_func(suffix=self.name)
-        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
-
-        @cache.dynamic_func(suffix=self.name)
-        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),
-                nabla_weight * grad_transform,
-            )
-
-        return eval_nabla_test_inner
-
-    def _make_eval_div_inner(self):
-        if not self.divergence_valid():
-            return None
-
-        @cache.dynamic_func(suffix=self.name)
-        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),
-                nabla_weight * grad_transform,
-            )
-
-        return eval_div_test_inner
-
-    def _make_eval_outer(self):
-        @cache.dynamic_func(suffix=self.name)
-        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
-
-        @cache.dynamic_func(suffix=self.name)
-        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),
-                nabla_weight * grad_transform,
-            )
-
-        return eval_nabla_test_outer
-
-    def _make_eval_div_outer(self):
-        if not self.divergence_valid():
-            return None
-
-        @cache.dynamic_func(suffix=self.name)
-        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),
-                nabla_weight * grad_transform,
-            )
-
-        return eval_div_test_outer
-
-    def _make_at_node(self):
-        @cache.dynamic_func(suffix=self.name)
-        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

warp/fem/field/trial.py

@@ -1,183 +0,0 @@
-import warp as wp
-from warp.fem import cache, utils
-from warp.fem.domain import GeometryDomain
-from warp.fem.space import FunctionSpace, SpacePartition
-from warp.fem.types import Sample, get_node_index_in_element
-
-from .field import SpaceField
-
-
-class TrialField(SpaceField):
-    """Field defined over a domain that can be used as a trial function"""
-
-    def __init__(
-        self,
-        space: FunctionSpace,
-        space_partition: SpacePartition,
-        domain: GeometryDomain,
-    ):
-        if domain.dimension == space.dimension - 1:
-            space = space.trace()
-
-        if domain.dimension != space.dimension:
-            raise ValueError("Incompatible space and domain dimensions")
-
-        if not space.topology.is_derived_from(space_partition.space_topology):
-            raise ValueError("Incompatible space and space partition topologies")
-
-        super().__init__(space, space_partition)
-
-        self.domain = 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 partition_node_count(self) -> int:
-        """Returns the number of nodes in the associated space topology partition"""
-        return self.space_partition.node_count()
-
-    @property
-    def name(self) -> str:
-        return self.space.name + "Trial"
-
-    def eval_arg_value(self, device) -> wp.codegen.StructInstance:
-        return self.space.space_arg_value(device)
-
-    def _make_element_eval_arg(self):
-        @cache.dynamic_struct(suffix=self.name)
-        class ElementEvalArg:
-            elt_arg: self.domain.ElementArg
-            eval_arg: self.EvalArg
-
-        return ElementEvalArg
-
-    def _make_eval_inner(self):
-        @cache.dynamic_func(suffix=self.name)
-        def eval_trial_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.trial_dof),
-            )
-            return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
-
-        return eval_trial_inner
-
-    def _make_eval_grad_inner(self):
-        if not self.gradient_valid():
-            return None
-
-        @cache.dynamic_func(suffix=self.name)
-        def eval_nabla_trial_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.trial_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.trial_dof),
-                nabla_weight * grad_transform,
-            )
-
-        return eval_nabla_trial_inner
-
-    def _make_eval_div_inner(self):
-        if not self.divergence_valid():
-            return None
-
-        @cache.dynamic_func(suffix=self.name)
-        def eval_div_trial_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.trial_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.trial_dof),
-                nabla_weight * grad_transform,
-            )
-
-        return eval_div_trial_inner
-
-    def _make_eval_outer(self):
-        @cache.dynamic_func(suffix=self.name)
-        def eval_trial_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.trial_dof),
-            )
-            return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
-
-        return eval_trial_outer
-
-    def _make_eval_grad_outer(self):
-        if not self.gradient_valid():
-            return None
-
-        @cache.dynamic_func(suffix=self.name)
-        def eval_nabla_trial_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.trial_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.trial_dof),
-                nabla_weight * grad_transform,
-            )
-
-        return eval_nabla_trial_outer
-
-    def _make_eval_div_outer(self):
-        if not self.divergence_valid():
-            return None
-
-        @cache.dynamic_func(suffix=self.name)
-        def eval_div_trial_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.trial_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.trial_dof),
-                nabla_weight * grad_transform,
-            )
-
-        return eval_div_trial_outer
-
-    def _make_at_node(self):
-        @cache.dynamic_func(suffix=self.name)
-        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.trial_dof)
-            )
-            return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
-
-        return at_node

warp/fem/space/collocated_function_space.py

@@ -1,102 +0,0 @@
-from typing import Optional
-
-import warp as wp
-from warp.fem import cache, utils
-from warp.fem.types import DofIndex, get_node_coord
-
-from .basis_space import BasisSpace
-from .dof_mapper import DofMapper, IdentityMapper
-from .function_space import FunctionSpace
-from .partition import SpacePartition, make_space_partition
-
-
-class CollocatedFunctionSpace(FunctionSpace):
-    """Function space where values are collocated at nodes"""
-
-    def __init__(self, basis: BasisSpace, dtype: type = float, dof_mapper: DofMapper = None):
-        super().__init__(topology=basis.topology)
-
-        self.dof_mapper = IdentityMapper(dtype) if dof_mapper is None else dof_mapper
-        self.dtype = self.dof_mapper.value_dtype
-        self.dof_dtype = self.dof_mapper.dof_dtype
-        self.VALUE_DOF_COUNT = self.dof_mapper.DOF_SIZE
-
-        self._basis = basis
-        self.SpaceArg = self._basis.BasisArg
-
-        self.ORDER = self._basis.ORDER
-
-        self.unit_dof_value = self._make_unit_dof_value(self.dof_mapper)
-
-        self.node_coords_in_element = self._basis.make_node_coords_in_element()
-        self.node_quadrature_weight = self._basis.make_node_quadrature_weight()
-        self.element_inner_weight = self._basis.make_element_inner_weight()
-        self.element_inner_weight_gradient = self._basis.make_element_inner_weight_gradient()
-        self.element_outer_weight = self._basis.make_element_outer_weight()
-        self.element_outer_weight_gradient = self._basis.make_element_outer_weight_gradient()
-
-        # For backward compatibility
-        if hasattr(basis, "node_grid"):
-            self.node_grid = basis.node_grid
-        if hasattr(basis, "node_triangulation"):
-            self.node_triangulation = basis.node_triangulation
-        if hasattr(basis, "node_tets"):
-            self.node_tets = basis.node_tets
-        if hasattr(basis, "node_hexes"):
-            self.node_hexes = basis.node_hexes
-        if hasattr(basis, "vtk_cells"):
-            self.vtk_cells = basis.vtk_cells
-
-    def space_arg_value(self, device):
-        return self._basis.basis_arg_value(device)
-
-    @property
-    def name(self):
-        return f"{self._basis.name}_{self.dof_mapper}".replace(".", "_")
-
-    @property
-    def degree(self):
-        """Maximum polynomial degree of the underlying basis"""
-        return self.ORDER
-
-    def make_field(
-        self,
-        space_partition: Optional[SpacePartition] = None,
-    ) -> "wp.fem.field.NodalField":
-        from warp.fem.field import NodalField
-
-        if space_partition is None:
-            space_partition = make_space_partition(space_topology=self.topology)
-
-        return NodalField(space=self, space_partition=space_partition)
-
-    def _make_unit_dof_value(self, dof_mapper: DofMapper):
-        @cache.dynamic_func(suffix=self.name)
-        def unit_dof_value(geo_arg: self.topology.ElementArg, space_arg: self.SpaceArg, dof: DofIndex):
-            return dof_mapper.dof_to_value(utils.unit_element(dof_mapper.dof_dtype(0.0), get_node_coord(dof)))
-
-        return unit_dof_value
-
-    def node_count(self):
-        return self.topology.node_count()
-
-    def node_positions(self, out: Optional[wp.array] = None) -> wp.array:
-        return self._basis.node_positions(out=out)
-
-    def trace(self) -> "CollocatedFunctionSpace":
-        return CollocatedFunctionSpaceTrace(self)
-
-
-class CollocatedFunctionSpaceTrace(CollocatedFunctionSpace):
-    """Trace of a :class:`CollocatedFunctionSpace`"""
-
-    def __init__(self, space: CollocatedFunctionSpace):
-        self._space = space
-        super().__init__(space._basis.trace(), space.dtype, space.dof_mapper)
-
-    @property
-    def name(self):
-        return f"{self._space.name}_Trace"
-
-    def __eq__(self, other: "CollocatedFunctionSpaceTrace") -> bool:
-        return self._space == other._space