Add SampleTriangleMeshBvh operator

Author: rycwo

metafold/func.py

@@ -45,6 +45,7 @@
     ParametrizationAsset,
     LineNetworkAsset,
     LineNetworkBvhAsset,
+    MeshBvhAsset,
 )
 
 if TYPE_CHECKING:
@@ -684,6 +685,44 @@ def __call__(self, eval_: Evaluator) -> TypedResult[Literal[FuncType.FLOAT]]:
         return cast(TypedResult[Literal[FuncType.FLOAT]], r)
 
 
+class SampleTriangleMeshBvh_Parameters(TypedDict, total=False):
+    xform: Mat4f
+
+
+class SampleTriangleMeshBvh(TypedFunc[Literal[FuncType.FLOAT]]):
+    def __init__(
+        self,
+        points: TypedFunc[Literal[FuncType.VEC3F]],
+        bvh_data: MeshBvhAsset,
+        mesh_data: TriangleMeshAsset,
+        parameters: Optional[SampleTriangleMeshBvh_Parameters] = None,
+    ):
+        self.inputs: Optional[dict[str, Func]]
+        self.inputs = {
+            "Points": points,
+        }
+        self.assets: Optional[Assets]
+        self.assets = {
+            "bvh_data": bvh_data,
+            "mesh_data": mesh_data,
+        }
+        self.parameters = parameters
+
+    @cache
+    def __call__(self, eval_: Evaluator) -> TypedResult[Literal[FuncType.FLOAT]]:
+        inputs: Optional[Inputs] = None
+        if self.inputs:
+            inputs = dict((k, v(eval_)) for k, v in self.inputs.items())
+        r = eval_(
+            "SampleTriangleMeshBvh",
+            inputs=inputs,
+            assets=self.assets,
+            # https://github.com/python/mypy/issues/4976#issuecomment-460971843
+            parameters=cast(Optional[Params], self.parameters),
+        )
+        return cast(TypedResult[Literal[FuncType.FLOAT]], r)
+
+
 class SampleVolume_Parameters(TypedDict, total=False):
     volume_offset: Vec3f
     volume_size: Vec3f

metafold/func.py.in

@@ -45,6 +45,7 @@ from .func_types import (
     ParametrizationAsset,
     LineNetworkAsset,
     LineNetworkBvhAsset,
+    MeshBvhAsset,
 )
 
 if TYPE_CHECKING:

metafold/func_types.py

@@ -87,6 +87,7 @@ class CustomShapeAsset(Asset):
 class ParametrizationAsset(Asset): ...
 class LineNetworkAsset(Asset): ...
 class LineNetworkBvhAsset(Asset): ...
+class MeshBvhAsset(Asset): ...
 
 
 class FuncType(Enum):