Spaces:
Sleeping
Sleeping
| from typing import Any, Optional | |
| import warp as wp | |
| from warp.fem.types import DofIndex, get_node_coord | |
| from warp.fem.geometry import GeometryPartition | |
| from warp.fem import cache, utils | |
| from .function_space import FunctionSpace | |
| from .dof_mapper import DofMapper, IdentityMapper | |
| from .partition import make_space_partition, SpacePartition | |
| from .basis_space import BasisSpace | |
| 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 | |
| def space_arg_value(self, device): | |
| return self._basis.basis_arg_value(device) | |
| def name(self): | |
| return f"{self._basis.name}_{self.dof_mapper}".replace(".", "_") | |
| 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): | |
| 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) | |
| def name(self): | |
| return f"{self._space.name}_Trace" | |
| def __eq__(self, other: "CollocatedFunctionSpaceTrace") -> bool: | |
| return self._space == other._space | |