NvidiaWarp-GarmentCode/warp/fem/geometry/deformed_geometry.py

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