diff --git a/docs/api/physicsnemo.sym.geometry.rst b/docs/api/physicsnemo.sym.geometry.rst index 476e0daf..d7cc081f 100644 --- a/docs/api/physicsnemo.sym.geometry.rst +++ b/docs/api/physicsnemo.sym.geometry.rst @@ -64,9 +64,16 @@ geometry.tessellation :members: :show-inheritance: +geometry.tessellation_warp +------------------------------------------------- + +.. automodule:: physicsnemo.sym.geometry.tessellation_warp + :members: + :show-inheritance: + geometry.geometry_dataloader ----------------------------- .. automodule:: physicsnemo.sym.geometry.geometry_dataloader :members: - :show-inheritance: \ No newline at end of file + :show-inheritance: diff --git a/docs/user_guide/features/csg_and_tessellated_module.rst b/docs/user_guide/features/csg_and_tessellated_module.rst index 7982e17f..d4713e50 100644 --- a/docs/user_guide/features/csg_and_tessellated_module.rst +++ b/docs/user_guide/features/csg_and_tessellated_module.rst @@ -6,22 +6,22 @@ Geometry Modules Constructive Solid Geometry --------------------------- -PhysicsNeMo Sym provides several 1D, 2D and 3D primitives that can be used for sampling point clouds required for the physics-informed training. +PhysicsNeMo Sym provides several 1D, 2D and 3D primitives that can be used for sampling point clouds required for the physics-informed training. These primitives support the standard boolean operations* like Union (``+`` ), Intersection (``&`` ) and Subtraction (``-`` ). The boolean -operations work on the signed distance fields of the differnt primitives. +operations work on the signed distance fields of the differnt primitives. -Below example shows a simple CSG primitive being built using PhysicsNeMo Sym. +Below example shows a simple CSG primitive being built using PhysicsNeMo Sym. .. code-block:: python :caption: Constructive solid geometry - + import numpy as np from physicsnemo.sym.geometry.primitives_3d import Box, Sphere, Cylinder from physicsnemo.sym.utils.io.vtk import var_to_polyvtk - + # number of points to sample nr_points = 100000 - + # make standard constructive solid geometry example # make primitives box = Box(point_1=(-1, -1, -1), point_2=(1, 1, 1)) @@ -29,12 +29,12 @@ Below example shows a simple CSG primitive being built using PhysicsNeMo Sym. cylinder_1 = Cylinder(center=(0, 0, 0), radius=0.5, height=2) cylinder_2 = cylinder_1.rotate(angle=float(np.pi / 2.0), axis="x") cylinder_3 = cylinder_1.rotate(angle=float(np.pi / 2.0), axis="y") - + # combine with boolean operations all_cylinders = cylinder_1 + cylinder_2 + cylinder_3 box_minus_sphere = box & sphere geo = box_minus_sphere - all_cylinders - + # sample geometry for plotting in Paraview s = geo.sample_boundary(nr_points=nr_points) var_to_polyvtk(s, "boundary") @@ -42,7 +42,7 @@ Below example shows a simple CSG primitive being built using PhysicsNeMo Sym. s = geo.sample_interior(nr_points=nr_points, compute_sdf_derivatives=True) var_to_polyvtk(s, "interior") print("Volume: {:.3f}".format(np.sum(s["area"]))) - + .. figure:: ../../images/user_guide/csg_demo.png :alt: Constructive Solid Geometry using PhysicsNeMo Sym primitives @@ -51,10 +51,10 @@ Below example shows a simple CSG primitive being built using PhysicsNeMo Sym. Constructive Solid Geometry using PhysicsNeMo Sym primitives -A complete list of primitives can be referred in ``physicsnemo.geometry.primitives_*`` +A complete list of primitives can be referred in ``physicsnemo.geometry.primitives_*`` .. note:: - While generating a complex primitive, it should be noted that the boolean operation are performed at the final stage, meaning it is invariant to + While generating a complex primitive, it should be noted that the boolean operation are performed at the final stage, meaning it is invariant to the order of the boolean operations. In other words, if you are subtracting a primitive from another primitive, and if you decide to add a different primitive in the area that is subtracted, you will not see the newly added primitive because the subtracted primitive. @@ -68,7 +68,7 @@ The geometry objects created also support operations like ``translate``, ``scale geo = geo.rotate(angle=np.pi / 4, axis="z") geo = geo.rotate(angle=np.pi / 4, axis="y") geo = geo.repeat(spacing=4.0, repeat_lower=(-1, -1, -1), repeat_higher=(1, 1, 1)) - + # sample geometry for plotting in Paraview s = geo.sample_boundary(nr_points=nr_points) var_to_polyvtk(s, "repeated_boundary") @@ -93,7 +93,7 @@ The CSG objects can be easily parameterized using sympy. An example of this is u from physicsnemo.sym.geometry.primitives_2d import Rectangle, Circle from physicsnemo.sym.utils.io.vtk import var_to_polyvtk from physicsnemo.sym.geometry.parameterization import Parameterization, Parameter - + # make plate with parameterized hole # make parameterized primitives plate = Rectangle(point_1=(-1, -1), point_2=(1, 1)) @@ -101,13 +101,13 @@ The CSG objects can be easily parameterized using sympy. An example of this is u parameterization = Parameterization({y_pos: (-1, 1)}) circle = Circle(center=(0, y_pos), radius=0.3, parameterization=parameterization) geo = plate - circle - + # sample geometry over entire parameter range s = geo.sample_boundary(nr_points=100000) var_to_polyvtk(s, "parameterized_boundary") s = geo.sample_interior(nr_points=100000) var_to_polyvtk(s, "parameterized_interior") - + # sample specific parameter s = geo.sample_boundary( nr_points=100000, parameterization=Parameterization({y_pos: 0}) @@ -117,7 +117,7 @@ The CSG objects can be easily parameterized using sympy. An example of this is u nr_points=100000, parameterization=Parameterization({y_pos: 0}) ) var_to_polyvtk(s, "y_pos_zero_interior") - + .. figure:: ../../images/user_guide/csg_parameterized_demo.png :alt: Parameterized Constructive Solid Geometry using PhysicsNeMo Sym primitives @@ -155,25 +155,25 @@ Defining Custom Primitives -------------------------- If you don't find a primitive defined for your application, it is easy to setup using the base classes from PhysicsNeMo Sym. All you need to do is come up with and appropriate -expression for the signed distance field and the surfaces of the geometry. An example is shown below. +expression for the signed distance field and the surfaces of the geometry. An example is shown below. .. code-block:: python :caption: Custom Primitive from sympy import Symbol, pi, sin, cos, sqrt, Min, Max, Abs - + from physicsnemo.sym.geometry.geometry import Geometry, csg_curve_naming from physicsnemo.sym.geometry.helper import _sympy_sdf_to_sdf from physicsnemo.sym.geometry.curve import SympyCurve, Curve from physicsnemo.sym.geometry.parameterization import Parameterization, Parameter, Bounds from physicsnemo.sym.geometry.primitives_3d import Cylinder from physicsnemo.sym.utils.io.vtk import var_to_polyvtk - + class InfiniteCylinder(Geometry): """ 3D Infinite Cylinder Axis parallel to z-axis, no caps on ends - + Parameters ---------- center : tuple with 3 ints or floats @@ -185,13 +185,13 @@ expression for the signed distance field and the surfaces of the geometry. An ex parameterization : Parameterization Parameterization of geometry. """ - + def __init__(self, center, radius, height, parameterization=Parameterization()): # make sympy symbols to use x, y, z = Symbol("x"), Symbol("y"), Symbol("z") h, r = Symbol(csg_curve_naming(0)), Symbol(csg_curve_naming(1)) theta = Symbol(csg_curve_naming(2)) - + # surface of the cylinder curve_parameterization = Parameterization( {h: (-1, 1), r: (0, 1), theta: (0, 2 * pi)} @@ -212,11 +212,11 @@ expression for the signed distance field and the surfaces of the geometry. An ex area=height * 2 * pi * radius, ) curves = [curve_1] - + # calculate SDF r_dist = sqrt((x - center[0]) ** 2 + (y - center[1]) ** 2) sdf = radius - r_dist - + # calculate bounds bounds = Bounds( { @@ -226,7 +226,7 @@ expression for the signed distance field and the surfaces of the geometry. An ex }, parameterization=parameterization, ) - + # initialize Infinite Cylinder super().__init__( curves, @@ -235,19 +235,19 @@ expression for the signed distance field and the surfaces of the geometry. An ex bounds=bounds, parameterization=parameterization, ) - - + + nr_points = 100000 - + cylinder_1 = Cylinder(center=(0, 0, 0), radius=0.5, height=2) cylinder_2 = InfiniteCylinder(center=(0, 0, 0), radius=0.5, height=2) - + s = cylinder_1.sample_interior(nr_points=nr_points, compute_sdf_derivatives=True) var_to_polyvtk(s, "interior_cylinder") - + s = cylinder_2.sample_interior(nr_points=nr_points, compute_sdf_derivatives=True) var_to_polyvtk(s, "interior_infinite_cylinder") - + .. figure:: ../../images/user_guide/custom_primitive_demo.png :alt: Custom primitive in PhysicsNeMo Sym. The cylinders are sliced to visualize the interior SDF @@ -256,7 +256,7 @@ expression for the signed distance field and the surfaces of the geometry. An ex Custom primitive in PhysicsNeMo Sym. The cylinders are sliced to visualize the interior SDF - + Tesselated Geometry ------------------- @@ -264,7 +264,7 @@ For more complicated shapes, PhysicsNeMo Sym allows geometries to be imported in The module also gives surface normals of the geometry for surface sampling. Once the geometry is imported, the point cloud can be used for training. An example of this can be found in :ref:`stl`. -Tesselated geometries can also be combined with the primitives +Tesselated geometries can also be combined with the primitives .. code-block:: python :caption: Tesselated Geometry @@ -273,17 +273,17 @@ Tesselated geometries can also be combined with the primitives from physicsnemo.sym.geometry.tessellation import Tessellation from physicsnemo.sym.geometry.primitives_3d import Plane from physicsnemo.sym.utils.io.vtk import var_to_polyvtk - + # number of points to sample nr_points = 100000 - + # make tesselated geometry from stl file geo = Tessellation.from_stl("./stl_files/tessellated_example.stl") - + # tesselated geometries can be combined with primitives cut_plane = Plane((0, -1, -1), (0, 1, 1)) geo = geo & cut_plane - + # sample geometry for plotting in Paraview s = geo.sample_boundary(nr_points=nr_points) var_to_polyvtk(s, "tessellated_boundary") @@ -291,10 +291,10 @@ Tesselated geometries can also be combined with the primitives s = geo.sample_interior(nr_points=nr_points, compute_sdf_derivatives=True) var_to_polyvtk(s, "tessellated_interior") print("Repeated Volume: {:.3f}".format(np.sum(s["area"]))) - + .. figure:: ../../images/user_guide/stl_demo_1.png - :alt: Tesselated Geometry sampling using PhysicsNeMo Sym + :alt: Tesselated Geometry sampling using PhysicsNeMo Sym :width: 80.0% :align: center @@ -307,6 +307,49 @@ Tesselated geometries can also be combined with the primitives Tesselated Geometry sampling using PhysicsNeMo Sym: Stanford bunny +Tesselated Geometry with Warp +----------------------------- + +PhysicsNeMo Sym also provides a Warp-based ``Tessellation`` module. +This module provides significant speedups for surface sampling using ``sample_interior`` +over the original ``Tessellation`` module. It is a drop-in replacement for the +original ``Tessellation`` module and can be used in the same way. + +.. code-block:: python + :caption: Tesselated Geometry with Warp + + from physicsnemo.sym.geometry.tessellation_warp import Tessellation + + # The rest of the code is the same. + +Below is a ``sample_boundary`` performance comparison of the original ``Tessellation`` module +and the Warp-based ``Tessellation`` module. + +.. list-table:: Tessellation Module Comparison (seconds, 100k sampled points) + :header-rows: 1 + + * - Mesh Size + - Original ``sample_boundary`` + - Warp-based ``sample_boundary`` + - Speedup + * - 10K + - 2.9698 + - 0.0367 + - 80x + * - 100K + - 20.6874 + - 0.1317 + - 150x + * - 1M + - 145.18 + - 0.4148 + - 350x + * - Usage + - ``from physicsnemo.sym.geometry.tessellation import Tessellation`` + - ``from physicsnemo.sym.geometry.tessellation_warp import Tessellation`` + - + + Signed Distance Fields (SDF) of Geometry objects ------------------------------------------------- @@ -320,7 +363,7 @@ applications. In physics-informed learning, it is also used to represent PhysicsNeMo' geometry module (CSG and Tesselation) computes the SDF (and its derivatives) on points sampled in the interior for use in the training pipelines. Additionally, the SDF can be computed on custom points using the ``.sdf`` attribute. - + PhysicsNeMo also provides a utility to recover the STL geometry from the SDF using marching cubes algorithm. For more details refer `here `_. @@ -352,14 +395,14 @@ Below example shows the use of these utilities for a CSG geometry. z = np.linspace(-2, 2, 100) xx, yy, zz = np.meshgrid(x, y, z, indexing='ij') - + # compute the SDF on the grid points sdf = geo.sdf( { - "x": xx.reshape(-1, 1), - "y": yy.reshape(-1, 1), + "x": xx.reshape(-1, 1), + "y": yy.reshape(-1, 1), "z": zz.reshape(-1, 1), - }, + }, params={} )["sdf"] @@ -400,17 +443,17 @@ Below example shows the use of these utilities for a Tesselation geometry. z = np.linspace(bounds["z"][0] - 1, bounds["z"][1] + 1, 100) xx, yy, zz = np.meshgrid(x, y, z, indexing='ij') - + # compute the SDF on the grid points sdf = geo.sdf( { - "x": xx.reshape(-1, 1), - "y": yy.reshape(-1, 1), + "x": xx.reshape(-1, 1), + "y": yy.reshape(-1, 1), "z": zz.reshape(-1, 1), - }, + }, params={} )["sdf"] - + # reconstruct the STL from SDF sdf_to_stl( sdf.reshape(100, 100, 100), # sdf field in [nx, ny, nz] shape diff --git a/physicsnemo/sym/geometry/tessellation_warp.py b/physicsnemo/sym/geometry/tessellation_warp.py new file mode 100644 index 00000000..2f4fdd70 --- /dev/null +++ b/physicsnemo/sym/geometry/tessellation_warp.py @@ -0,0 +1,422 @@ +# 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 + + +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)) + + 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 + + +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, + ) + + @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) + + @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 diff --git a/test/test_geometry.py b/test/test_geometry.py index cef02de4..54663c46 100644 --- a/test/test_geometry.py +++ b/test/test_geometry.py @@ -36,6 +36,7 @@ Torus, ) from physicsnemo.sym.geometry.tessellation import Tessellation +from physicsnemo.sym.geometry.tessellation_warp import Tessellation as TessellationWarp from physicsnemo.sym.utils.io.vtk import var_to_polyvtk dir_path = Path(__file__).parent @@ -238,4 +239,87 @@ def interior_criteria(invar, params): ) -test_primitives() +def test_tessellation_equivalency(): + """ + Test equivalency between original Tessellation and Warp Tessellation implementations. + + This test compares the sample_boundary method from both implementations + to ensure they produce equivalent results. + It roughly reproduces relevant checks from test_primitives. + """ + + # Set up test parameters. + nr_points = 1000 + stl_file = dir_path / "stls/cube.stl" + + # Create original and Warp Tessellation objects. + tess_original = Tessellation.from_stl(stl_file) + tess_warp = TessellationWarp.from_stl(stl_file) + + # Sample boundary. + boundary_original = tess_original.sample_boundary(nr_points) + boundary_warp = tess_warp.sample_boundary(nr_points) + + # Check that both dictionaries have the same keys. + assert set(boundary_original.keys()) == set(boundary_warp.keys()), ( + f"Different keys: original={set(boundary_original.keys())}, warp={set(boundary_warp.keys())}" + ) + + # Check that shapes are consistent. + for key in boundary_original.keys(): + assert boundary_original[key].shape == boundary_warp[key].shape, ( + f"Shape mismatch for key '{key}': original={boundary_original[key].shape}, warp={boundary_warp[key].shape}" + ) + + # Check that total surface areas are approximately equal. + total_area_original = np.sum(boundary_original["area"]) + total_area_warp = np.sum(boundary_warp["area"]) + assert np.isclose(total_area_original, total_area_warp, rtol=1e-2), ( + f"Total area mismatch: original={total_area_original}, warp={total_area_warp}" + ) + + # Verify the expected surface area for a unit cube (should be 6.0). + assert np.isclose(total_area_original, 6.0, rtol=1e-1), ( + f"Expected cube surface area ~6.0, got {total_area_original}" + ) + assert np.isclose(total_area_warp, 6.0, rtol=1e-1), ( + f"Expected cube surface area ~6.0, got {total_area_warp}" + ) + + # Check that normals are unit vectors (approximately). + for impl_name, boundary in ( + ("original", boundary_original), + ("warp", boundary_warp), + ): + normal_magnitudes = np.sqrt( + boundary["normal_x"] ** 2 + + boundary["normal_y"] ** 2 + + boundary["normal_z"] ** 2 + ) + assert np.allclose(normal_magnitudes, 1.0, rtol=1e-3), ( + f"Normal vectors are not unit vectors in {impl_name} implementation" + ) + + # Check that points are on the cube surface (SDF should be close to 0). + for impl_name, boundary in ( + ("original", boundary_original), + ("warp", boundary_warp), + ): + # Create points array for SDF computation. + points_dict = {"x": boundary["x"], "y": boundary["y"], "z": boundary["z"]} + + # Compute SDF using original tessellation. + sdf_result = tess_original.sdf(points_dict, {}) + sdf_values = sdf_result["sdf"] + + # All boundary points should have SDF close to 0. + assert np.allclose(sdf_values, 0.0, atol=1e-3), ( + f"Boundary points not on surface in {impl_name} implementation, max SDF: {np.max(np.abs(sdf_values))}" + ) + + print("Tessellation equivalency test passed!") + + +if __name__ == "__main__": + test_primitives() + test_tessellation_equivalency()