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() @property def name(self): return f"DefGeo_{self.field.name}" # Geometry device interface @cache.cached_arg_value 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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()) @cache.dynamic_func(suffix=self.name) def cell_measure(args: self.CellArg, s: Sample): return wp.abs(wp.determinant(self.cell_deformation_gradient(args, s))) * REF_MEASURE return cell_measure @wp.func def cell_normal(args: Any, s: Sample): return wp.vec2(0.0) def _make_side_arg(self): @cache.dynamic_struct(suffix=self.name) class SideArg: base_arg: self.base.SideArg trace_arg: self.field_trace.EvalArg field_arg: self.field.EvalArg return SideArg @cache.cached_arg_value 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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)) @wp.func def _side_measure(F: wp.vec2): return wp.length(F) @wp.func 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) @wp.func def _side_normal(F: wp.vec2): return wp.normalize(wp.vec2(-F[1], F[0])) @wp.func 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()) @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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): @cache.dynamic_func(suffix=self.name) 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