Spaces:
Paused
Paused
| from typing import Any | |
| import warp as wp | |
| from .geometry import Geometry | |
| from warp.fem.types import Sample, ElementIndex, Coords, make_free_sample | |
| from warp.fem import cache | |
| _mat32 = wp.mat(shape=(3, 2), dtype=float) | |
| class DeformedGeometry(Geometry): | |
| def __init__(self, field): | |
| """Constructs a Deformed Geometry from a displacement field defined over a base geometry""" | |
| from warp.fem.field import DiscreteField | |
| self.field: DiscreteField = field | |
| self.base = self.field.space.geometry | |
| self.dimension = self.base.dimension | |
| if not wp.types.type_is_vector(field.dtype) or wp.types.type_length(field.dtype) != self.dimension: | |
| raise ValueError("Invalid value type for position field") | |
| self.CellArg = self.field.ElementEvalArg | |
| self.field_trace = field.trace() | |
| self.SideArg = self._make_side_arg() | |
| self.SideIndexArg = self.base.SideIndexArg | |
| self.cell_count = self.base.cell_count | |
| self.vertex_count = self.base.vertex_count | |
| self.side_count = self.base.side_count | |
| self.boundary_side_count = self.base.boundary_side_count | |
| self.reference_cell = self.base.reference_cell | |
| self.reference_side = self.base.reference_side | |
| self.side_index_arg_value = self.base.side_index_arg_value | |
| self.cell_position = self._make_cell_position() | |
| self.cell_deformation_gradient = self._make_cell_deformation_gradient() | |
| self.cell_inverse_deformation_gradient = self._make_cell_inverse_deformation_gradient() | |
| self.cell_measure = self._make_cell_measure() | |
| self.boundary_side_index = self.base.boundary_side_index | |
| self.side_to_cell_arg = self._make_side_to_cell_arg() | |
| self.side_position = self._make_side_position() | |
| self.side_deformation_gradient = self._make_side_deformation_gradient() | |
| self.side_inner_cell_index = self._make_side_inner_cell_index() | |
| self.side_outer_cell_index = self._make_side_outer_cell_index() | |
| self.side_inner_cell_coords = self._make_side_inner_cell_coords() | |
| self.side_outer_cell_coords = self._make_side_outer_cell_coords() | |
| self.side_from_cell_coords = self._make_side_from_cell_coords() | |
| self.side_inner_inverse_deformation_gradient = self._make_side_inner_inverse_deformation_gradient() | |
| self.side_outer_inverse_deformation_gradient = self._make_side_outer_inverse_deformation_gradient() | |
| self.side_measure = self._make_side_measure() | |
| self.side_measure_ratio = self._make_side_measure_ratio() | |
| self.side_normal = self._make_side_normal() | |
| def name(self): | |
| return f"DefGeo_{self.field.name}" | |
| # Geometry device interface | |
| def cell_arg_value(self, device) -> "DeformedGeometry.CellArg": | |
| args = self.CellArg() | |
| args.elt_arg = self.base.cell_arg_value(device) | |
| args.eval_arg = self.field.eval_arg_value(device) | |
| return args | |
| def _make_cell_position(self): | |
| def cell_position(cell_arg: self.CellArg, s: Sample): | |
| return self.field.eval_inner(cell_arg, s) + self.base.cell_position(cell_arg.elt_arg, s) | |
| return cell_position | |
| def _make_cell_deformation_gradient(self): | |
| def cell_deformation_gradient(cell_arg: self.CellArg, s: Sample): | |
| return self.field.eval_reference_grad_inner(cell_arg, s) + self.base.cell_deformation_gradient( | |
| cell_arg.elt_arg, s | |
| ) | |
| return cell_deformation_gradient | |
| def _make_cell_inverse_deformation_gradient(self): | |
| def cell_inverse_deformation_gradient(cell_arg: self.CellArg, s: Sample): | |
| return wp.inverse(self.cell_deformation_gradient(cell_arg, s)) | |
| return cell_inverse_deformation_gradient | |
| def _make_cell_measure(self): | |
| REF_MEASURE = wp.constant(self.reference_cell().measure()) | |
| def cell_measure(args: self.CellArg, s: Sample): | |
| return wp.abs(wp.determinant(self.cell_deformation_gradient(args, s))) * REF_MEASURE | |
| return cell_measure | |
| def cell_normal(args: Any, s: Sample): | |
| return wp.vec2(0.0) | |
| def _make_side_arg(self): | |
| class SideArg: | |
| base_arg: self.base.SideArg | |
| trace_arg: self.field_trace.EvalArg | |
| field_arg: self.field.EvalArg | |
| return SideArg | |
| def side_arg_value(self, device) -> "DeformedGeometry.SideArg": | |
| args = self.SideArg() | |
| args.base_arg = self.base.side_arg_value(device) | |
| args.field_arg = self.field.eval_arg_value(device) | |
| args.trace_arg = self.field_trace.eval_arg_value(device) | |
| return args | |
| def _make_side_position(self): | |
| def side_position(args: self.SideArg, s: Sample): | |
| trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg) | |
| return self.field_trace.eval_inner(trace_arg, s) + self.base.side_position(args.base_arg, s) | |
| return side_position | |
| def _make_side_deformation_gradient(self): | |
| def side_deformation_gradient(args: self.SideArg, s: Sample): | |
| base_def_grad = self.base.side_deformation_gradient(args.base_arg, s) | |
| trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg) | |
| Du = self.field_trace.eval_grad_inner(trace_arg, s) | |
| return base_def_grad + Du * base_def_grad | |
| return side_deformation_gradient | |
| def _make_side_inner_inverse_deformation_gradient(self): | |
| def side_inner_inverse_deformation_gradient(args: self.SideArg, s: Sample): | |
| cell_index = self.side_inner_cell_index(args, s.element_index) | |
| cell_coords = self.side_inner_cell_coords(args, s.element_index, s.element_coords) | |
| cell_arg = self.side_to_cell_arg(args) | |
| return self.cell_inverse_deformation_gradient(cell_arg, make_free_sample(cell_index, cell_coords)) | |
| def _make_side_outer_inverse_deformation_gradient(self): | |
| def side_outer_inverse_deformation_gradient(args: self.SideArg, s: Sample): | |
| cell_index = self.side_outer_cell_index(args, s.element_index) | |
| cell_coords = self.side_outer_cell_coords(args, s.element_index, s.element_coords) | |
| cell_arg = self.side_to_cell_arg(args) | |
| return self.cell_inverse_deformation_gradient(cell_arg, make_free_sample(cell_index, cell_coords)) | |
| def _side_measure(F: wp.vec2): | |
| return wp.length(F) | |
| def _side_measure(F: _mat32): | |
| Fcross = wp.vec3( | |
| F[1, 0] * F[2, 1] - F[2, 0] * F[1, 1], | |
| F[2, 0] * F[0, 1] - F[0, 0] * F[2, 1], | |
| F[0, 0] * F[1, 1] - F[1, 0] * F[0, 1], | |
| ) | |
| return wp.length(Fcross) | |
| def _side_normal(F: wp.vec2): | |
| return wp.normalize(wp.vec2(-F[1], F[0])) | |
| def _side_normal(F: _mat32): | |
| Fcross = wp.vec3( | |
| F[1, 0] * F[2, 1] - F[2, 0] * F[1, 1], | |
| F[2, 0] * F[0, 1] - F[0, 0] * F[2, 1], | |
| F[0, 0] * F[1, 1] - F[1, 0] * F[0, 1], | |
| ) | |
| return wp.normalize(Fcross) | |
| def _make_side_measure(self): | |
| REF_MEASURE = wp.constant(self.reference_side().measure()) | |
| def side_measure(args: self.SideArg, s: Sample): | |
| F = self.side_deformation_gradient(args, s) | |
| return DeformedGeometry._side_measure(F) * REF_MEASURE | |
| return side_measure | |
| def _make_side_measure_ratio(self): | |
| def side_measure_ratio(args: self.SideArg, s: Sample): | |
| inner = self.side_inner_cell_index(args, s.element_index) | |
| outer = self.side_outer_cell_index(args, s.element_index) | |
| inner_coords = self.side_inner_cell_coords(args, s.element_index, s.element_coords) | |
| outer_coords = self.side_outer_cell_coords(args, s.element_index, s.element_coords) | |
| cell_arg = self.side_to_cell_arg(args) | |
| return self.side_measure(args, s) / wp.min( | |
| self.cell_measure(cell_arg, make_free_sample(inner, inner_coords)), | |
| self.cell_measure(cell_arg, make_free_sample(outer, outer_coords)), | |
| ) | |
| return side_measure_ratio | |
| def _make_side_normal(self): | |
| def side_normal(args: self.SideArg, s: Sample): | |
| F = self.side_deformation_gradient(args, s) | |
| return DeformedGeometry._side_normal(F) | |
| return side_normal | |
| def _make_side_inner_cell_index(self): | |
| def side_inner_cell_index(args: self.SideArg, side_index: ElementIndex): | |
| return self.base.side_inner_cell_index(args.base_arg, side_index) | |
| return side_inner_cell_index | |
| def _make_side_outer_cell_index(self): | |
| def side_outer_cell_index(args: self.SideArg, side_index: ElementIndex): | |
| return self.base.side_outer_cell_index(args.base_arg, side_index) | |
| return side_outer_cell_index | |
| def _make_side_inner_cell_coords(self): | |
| def side_inner_cell_coords(args: self.SideArg, side_index: ElementIndex, side_coords: Coords): | |
| return self.base.side_inner_cell_coords(args.base_arg, side_index, side_coords) | |
| return side_inner_cell_coords | |
| def _make_side_outer_cell_coords(self): | |
| def side_outer_cell_coords(args: self.SideArg, side_index: ElementIndex, side_coords: Coords): | |
| return self.base.side_outer_cell_coords(args.base_arg, side_index, side_coords) | |
| return side_outer_cell_coords | |
| def _make_side_from_cell_coords(self): | |
| def side_from_cell_coords( | |
| args: self.SideArg, | |
| side_index: ElementIndex, | |
| cell_index: ElementIndex, | |
| cell_coords: Coords, | |
| ): | |
| return self.base.side_from_cell_coords(args.base_arg, side_index, cell_index, cell_coords) | |
| return side_from_cell_coords | |
| def _make_side_to_cell_arg(self): | |
| def side_to_cell_arg(side_arg: self.SideArg): | |
| return self.CellArg(self.base.side_to_cell_arg(side_arg.base_arg), side_arg.field_arg) | |
| return side_to_cell_arg | |