PyRXMesh

PyRXMesh brings RXMesh GPU mesh processing to Python. PyRXMesh allows using Python and PyTorch to organize an application, while RXMesh accelerates the mesh processing and automatic differentiation on the GPU.

PyRXMesh provides:

• triangle meshes built from OBJ files or NumPy arrays
• vertex, edge, and face attributes on the host and GPU
• NumPy, SciPy, and PyTorch interoperability
• dense and sparse matrices and linear solvers
• small compiled plugins for custom RXMesh CUDA operations
• differentiable scalar energies that can be interop with PyTorch autograd.

Install

python -m pip install PyRXMesh

PyRXMesh needs an NVIDIA GPU and a driver compatible with the CUDA version in the wheel. PyTorch and SciPy are optional. Install SciPy normally, and choose the CUDA-enabled PyTorch build for your system from the PyTorch installation page.

Visualization is optional and it depends on Polyscope:

python -m pip install "PyRXMesh[viz]"

For a source build, see docs/DEVELOPING.md.

Create a mesh

Initialize RXMesh once, then load a mesh or build one from arrays:

import pyrxmesh as rx

rx.init(0) #0: GPU id

vertices, faces = rx.create_plane(32, 32, dx=0.05)
mesh = rx.RXMeshStatic(vertices, faces)

print(mesh.num_vertices, mesh.num_edges, mesh.num_faces)

velocity = mesh.add_vertex_attribute("velocity", dtype="float32", dim=3)
velocity.reset(0, "device")

RXMeshStatic("mesh.obj") loads an OBJ file. Mesh and attribute arrays use RXMesh linear order by default and maps are available when an application needs input/global order.

See examples/load_and_show.py for mesh inspection and Polyscope visualization, and the tests examples on mesh construction and export and attributes handling.

Automatic differentiation and PyTorch interop

A differentiable plugin defines one or more scalar energy terms as RXMesh CUDA lambdas. We provide a starter plugin/python module which creates an independent project with boilerplate code (pyproject.toml, CMakeLists.txt, etc) needed to compile the CUDA lambda function and build their python bindings:

python -m pyrxmesh.diff_plugin my_energy

The user writes their energy terms in .cu file. Once built, PyRXMesh evaluates its value and gradient on the GPU and exposes the result to PyTorch:

import torch
import my_energy

energy = my_energy.make_energy(mesh, {"weight": 0.5})
x = torch.as_tensor(
    mesh.vertices(), dtype=torch.float32, device="cuda"
).requires_grad_()

loss = energy.torch(x)
loss.backward()

The gradient flows through whatever PyTorch computation produced x. For an optimizer that accepts a value and gradient directly, use energy.value_and_grad() instead of creating an autograd node.

The starter contains the CUDA term, Python package, build files, and run script. See Writing Differentiable Energy Plugins for the complete workflow, including query selection, fixed attributes, row order, and zero-copy input.

Exchange data with Python

Attributes are attached to vertices, edges, or faces. They use SoA layout by default and may have host storage, device storage, or both.

coords = mesh.input_vertex_coordinates()

host_view = coords.to_numpy("host") #Zero-copy numpy view
device_tensor = coords.to_torch("device") #Zero-copy pytorch view 
owned_array = coords.to_numpy_copy(source="host") #deep copy 

Unsuffixed to_* methods return a zero-copy view when supported. Methods ending in _copy return independent storage. Host and device allocations are separate; use attribute.move(source, target) when one side must receive changes from the other.

Mesh ordering maps, face arrays, and the Polyscope edge permutation are available directly from RXMeshStatic. See tests/test_core_mesh.py and tests/test_attributes.py for the supported operations.

Add custom CUDA operations

When an operation cannot be expressed efficiently from Python, generate a small plugin and write only the RXMesh device lambda:

python -m pyrxmesh.plugin my_kernels

The generated package builds against the PyRXMesh installation in the active environment. See Writing Custom CUDA Plugins and examples/custom_kernel_plugin.

Matrices and solvers

PyRXMesh exposes dense matrices, CSR sparse matrices, mesh-derived sparsity patterns, iterative solvers, and optional direct solvers. Dtypes, storage locations, matrix order, and mesh query patterns use short string arguments where they apply.

Documentation

• Writing Differentiable Energy Plugins
• Writing Custom CUDA Plugins
• Developing PyRXMesh
• RXMesh CUDA/C++ documentation