# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from functools import partial
import numpy as np
import warp as wp
from stl import mesh as np_mesh
from .geometry import Geometry
from .parameterization import Parameterization, Bounds, Parameter
from .curve import Curve
from physicsnemo.sym.constants import diff_str
from physicsnemo.utils.sdf import signed_distance_field
@wp.kernel
def compute_triangle_areas_warp(
v0: wp.array(dtype=wp.vec3),
v1: wp.array(dtype=wp.vec3),
v2: wp.array(dtype=wp.vec3),
areas: wp.array(dtype=float),
):
"""Warp kernel to compute triangle areas."""
# This is a warp implementation of the _area_of_triangles function in tessellation.py.
tid = wp.tid()
# Get triangle vertices.
vertex0 = v0[tid]
vertex1 = v1[tid]
vertex2 = v2[tid]
# Compute edges.
edge1 = vertex1 - vertex0
edge2 = vertex2 - vertex0
# Cross product gives area vector.
cross = wp.cross(edge1, edge2)
areas[tid] = wp.length(cross) * 0.5
@wp.kernel
def sample_triangles_warp(
v0: wp.array(dtype=wp.vec3),
v1: wp.array(dtype=wp.vec3),
v2: wp.array(dtype=wp.vec3),
normals: wp.array(dtype=wp.vec3),
triangle_indices: wp.array(dtype=int),
random_r1: wp.array(dtype=float),
random_r2: wp.array(dtype=float),
points_x: wp.array(dtype=float),
points_y: wp.array(dtype=float),
points_z: wp.array(dtype=float),
normal_x: wp.array(dtype=float),
normal_y: wp.array(dtype=float),
normal_z: wp.array(dtype=float),
):
"""Warp kernel to sample points on triangles using barycentric coordinates."""
# This is a warp implementation of the _sample_triangle function in tessellation.py.
tid = wp.tid()
# Get triangle index for this point.
tri_idx = triangle_indices[tid]
# Get random values.
r1 = random_r1[tid]
r2 = random_r2[tid]
# Barycentric sampling (uniform distribution in triangle).
s1 = wp.sqrt(r1)
u = 1.0 - s1
v = (1.0 - r2) * s1
w = r2 * s1
# Get triangle vertices.
vertex0 = v0[tri_idx]
vertex1 = v1[tri_idx]
vertex2 = v2[tri_idx]
# Sample point using barycentric coordinates.
sampled_point = u * vertex0 + v * vertex1 + w * vertex2
# Store coordinates.
points_x[tid] = sampled_point[0]
points_y[tid] = sampled_point[1]
points_z[tid] = sampled_point[2]
# Get and normalize normal.
normal = normals[tri_idx]
normal_length = wp.length(normal)
if normal_length > 0.0:
normalized_normal = normal / normal_length
normal_x[tid] = normalized_normal[0]
normal_y[tid] = normalized_normal[1]
normal_z[tid] = normalized_normal[2]
else:
normal_x[tid] = 0.0
normal_y[tid] = 0.0
normal_z[tid] = 1.0
[docs]
class WarpTessellationSampler:
"""Warp-accelerated tessellation sampling.
Parameters
----------
mesh : Mesh (numpy-stl)
A mesh that defines the surface of the geometry.
device : str, optional
Device to use for Warp operations ('cuda' or 'cpu').
If None, automatically sets the device to the current device.
seed : int, optional
Seed for the random number generator. If None, uses default seeding.
"""
def __init__(self, mesh, device=None, seed=None):
self.num_triangles = len(mesh.v0)
self.rng = np.random.default_rng(seed)
self.device = device if device is not None else wp.get_device()
# Convert mesh data to Warp format.
v0_data = [(float(v[0]), float(v[1]), float(v[2])) for v in mesh.v0]
v1_data = [(float(v[0]), float(v[1]), float(v[2])) for v in mesh.v1]
v2_data = [(float(v[0]), float(v[1]), float(v[2])) for v in mesh.v2]
normals_data = [(float(n[0]), float(n[1]), float(n[2])) for n in mesh.normals]
self.v0_wp = wp.array(v0_data, dtype=wp.vec3, device=self.device)
self.v1_wp = wp.array(v1_data, dtype=wp.vec3, device=self.device)
self.v2_wp = wp.array(v2_data, dtype=wp.vec3, device=self.device)
self.normals_wp = wp.array(normals_data, dtype=wp.vec3, device=self.device)
# Pre-compute triangle areas.
self.areas_wp = wp.zeros(self.num_triangles, device=self.device)
wp.launch(
compute_triangle_areas_warp,
dim=self.num_triangles,
inputs=[self.v0_wp, self.v1_wp, self.v2_wp, self.areas_wp],
device=self.device,
)
# Copy areas to CPU for probability computation.
areas_host = self.areas_wp.numpy()
self.triangle_areas = areas_host
self.triangle_probabilities = areas_host / np.sum(areas_host)
self.total_area = float(np.sum(areas_host))
[docs]
def sample_warp(
self,
nr_points: int,
parameterization: Parameterization = None,
quasirandom: bool = False,
) -> tuple[dict, dict]:
"""
Warp-accelerated sampling function.
Parameters
----------
nr_points : int
Number of points to sample.
parameterization : Parameterization, optional
Parameterization of the geometry. Default is an empty Parameterization.
quasirandom : bool, optional
If True, use quasirandom sampling. Default is False.
Returns
-------
invar : dict
Dictionary containing the sampled points and normals.
params : dict
Dictionary containing the parameters of the sampled points.
"""
nr_points = int(nr_points)
# Step 1: Distribute points across triangles based on area.
triangle_indices = self.rng.choice(
self.num_triangles, size=nr_points, p=self.triangle_probabilities
)
# Step 2: Generate random numbers for barycentric sampling.
r1 = self.rng.uniform(0, 1, size=nr_points).astype(np.float32)
r2 = self.rng.uniform(0, 1, size=nr_points).astype(np.float32)
# Step 3: Create Warp arrays.
triangle_indices_wp = wp.array(triangle_indices, dtype=int, device=self.device)
r1_wp = wp.array(r1, dtype=float, device=self.device)
r2_wp = wp.array(r2, dtype=float, device=self.device)
# Step 4: Create output arrays.
points_x_wp = wp.zeros(nr_points, device=self.device)
points_y_wp = wp.zeros(nr_points, device=self.device)
points_z_wp = wp.zeros(nr_points, device=self.device)
normal_x_wp = wp.zeros(nr_points, device=self.device)
normal_y_wp = wp.zeros(nr_points, device=self.device)
normal_z_wp = wp.zeros(nr_points, device=self.device)
# Step 5: Launch sampling kernel.
wp.launch(
sample_triangles_warp,
dim=nr_points,
inputs=[
self.v0_wp,
self.v1_wp,
self.v2_wp,
self.normals_wp,
triangle_indices_wp,
r1_wp,
r2_wp,
points_x_wp,
points_y_wp,
points_z_wp,
normal_x_wp,
normal_y_wp,
normal_z_wp,
],
device=self.device,
)
# Step 6: Copy results back to CPU.
x = points_x_wp.numpy().reshape(-1, 1)
y = points_y_wp.numpy().reshape(-1, 1)
z = points_z_wp.numpy().reshape(-1, 1)
nx = normal_x_wp.numpy().reshape(-1, 1)
ny = normal_y_wp.numpy().reshape(-1, 1)
nz = normal_z_wp.numpy().reshape(-1, 1)
# Step 7: Compute area weights (matching original implementation).
area_weight = self.total_area / nr_points
area = np.full((nr_points, 1), area_weight, dtype=np.float32)
# Step 8: Create invar dictionary matching original format.
invar = {
"x": x,
"y": y,
"z": z,
"normal_x": nx,
"normal_y": ny,
"normal_z": nz,
"area": area,
}
# Step 9: Handle parameterization.
if parameterization is None:
parameterization = Parameterization()
params = parameterization.sample(nr_points, quasirandom=quasirandom)
return invar, params
[docs]
class Tessellation(Geometry):
"""
Tessellation is a geometry that uses Warp to accelerate the sampling of a tessellated geometry.
Parameters
----------
mesh : Mesh (numpy-stl)
A mesh that defines the surface of the geometry.
airtight : bool
If the geometry is airtight or not. If false sample everywhere for interior.
parameterization : Parameterization, optional
Parameterization of the geometry. Default is an empty Parameterization.
device : str, optional
Device to use for Warp operations ('cuda' or 'cpu'). If None, automatically
detects the best available device.
seed : int, optional
Seed for the random number generator. If None, uses default seeding.
"""
def __init__(
self,
mesh,
airtight: bool = True,
parameterization: Parameterization = Parameterization(),
device: str = None,
seed: int = None,
):
# Create curve with Warp-accelerated sampling.
curves = [
Curve(
WarpTessellationSampler(mesh, device=device, seed=seed).sample_warp,
dims=3,
parameterization=parameterization,
)
]
bounds = Bounds(
{
Parameter("x"): (
float(np.min(mesh.vectors[:, :, 0])),
float(np.max(mesh.vectors[:, :, 0])),
),
Parameter("y"): (
float(np.min(mesh.vectors[:, :, 1])),
float(np.max(mesh.vectors[:, :, 1])),
),
Parameter("z"): (
float(np.min(mesh.vectors[:, :, 2])),
float(np.max(mesh.vectors[:, :, 2])),
),
},
parameterization=parameterization,
)
super().__init__(
curves,
partial(Tessellation.sdf, mesh.vectors, airtight),
dims=3,
bounds=bounds,
parameterization=parameterization,
)
[docs]
@classmethod
def from_stl(
cls,
filename,
airtight=True,
parameterization=Parameterization(),
device=None,
seed=None,
):
"""
Create a Tessellation geometry from an STL file.
Parameters
----------
filename : str
Path to the STL mesh file.
airtight : bool, optional
If True, assumes the geometry is airtight. Default is True.
parameterization : Parameterization, optional
Parameterization of the geometry. Default is an empty Parameterization.
device : str, optional
Device to use for Warp operations ('cuda' or 'cpu'). If None, automatically
detects the best available device.
seed : int, optional
Seed for the random number generator. If None, uses default seeding.
Returns
-------
Tessellation
An instance of Tessellation initialized with the mesh from the STL file.
"""
# Read in mesh.
mesh = np_mesh.Mesh.from_file(filename)
return cls(mesh, airtight, parameterization, device, seed)
[docs]
@staticmethod
def sdf(triangles, airtight, invar, params, compute_sdf_derivatives=False):
"""Simple copy of the sdf function in tessellation.py."""
points = np.stack([invar["x"], invar["y"], invar["z"]], axis=1)
minx, maxx, miny, maxy, minz, maxz = Tessellation.find_mins_maxs(points)
max_dis = max(max((maxx - minx), (maxy - miny)), (maxz - minz))
store_triangles = np.array(triangles, dtype=np.float64)
store_triangles[:, :, 0] -= minx
store_triangles[:, :, 1] -= miny
store_triangles[:, :, 2] -= minz
store_triangles *= 1 / max_dis
store_triangles = store_triangles.reshape(-1, 3)
points[:, 0] -= minx
points[:, 1] -= miny
points[:, 2] -= minz
points *= 1 / max_dis
points = points.astype(np.float64).flatten()
outputs = {}
if airtight:
sdf_field, sdf_derivative = signed_distance_field(
store_triangles,
np.arange((store_triangles.shape[0])),
points,
include_hit_points=True,
)
if isinstance(sdf_field, wp.types.array):
sdf_field = sdf_field.numpy()
if isinstance(sdf_derivative, wp.types.array):
sdf_derivative = sdf_derivative.numpy()
sdf_field = -np.expand_dims(max_dis * sdf_field, axis=1)
sdf_derivative = sdf_derivative.reshape(-1)
else:
sdf_field = np.zeros_like(invar["x"])
outputs["sdf"] = sdf_field
if compute_sdf_derivatives:
sdf_derivative = -(sdf_derivative - points)
sdf_derivative = np.reshape(
sdf_derivative, (sdf_derivative.shape[0] // 3, 3)
)
sdf_derivative = sdf_derivative / np.linalg.norm(
sdf_derivative, axis=1, keepdims=True
)
outputs["sdf" + diff_str + "x"] = sdf_derivative[:, 0:1]
outputs["sdf" + diff_str + "y"] = sdf_derivative[:, 1:2]
outputs["sdf" + diff_str + "z"] = sdf_derivative[:, 2:3]
return outputs
@staticmethod
def find_mins_maxs(points):
minx = float(np.min(points[:, 0]))
miny = float(np.min(points[:, 1]))
minz = float(np.min(points[:, 2]))
maxx = float(np.max(points[:, 0]))
maxy = float(np.max(points[:, 1]))
maxz = float(np.max(points[:, 2]))
return minx, maxx, miny, maxy, minz, maxz