Uniform mesh sampling (#14071)

# Objective

Allow random sampling from the surfaces of triangle meshes.

## Solution

This has two parts.

Firstly, rendering meshes can now yield their collections of triangles
through a method `Mesh::triangles`. This has signature
```rust
pub fn triangles(&self) -> Result<Vec<Triangle3d>, MeshTrianglesError> { //... }
```

and fails in a variety of cases — the most obvious of these is that the
mesh must have either the `TriangleList` or `TriangleStrip` topology,
and the others correspond to malformed vertex or triangle-index data.

With that in hand, we have the second piece, which is
`UniformMeshSampler`, which is a `Vec3`-valued
[distribution](https://docs.rs/rand/latest/rand/distributions/trait.Distribution.html)
that samples uniformly from collections of triangles. It caches the
triangles' distribution of areas so that after its initial setup,
sampling is allocation-free. It is constructed via
`UniformMeshSampler::try_new`, which looks like this:
```rust
pub fn try_new<T: Into<Vec<Triangle3d>>>(triangles: T) -> Result<Self, ZeroAreaMeshError> { //... }
```

It fails if the collection of triangles has zero area. 

The sum of these parts means that you can sample random points from a
mesh as follows:
```rust
let triangles = my_mesh.triangles().unwrap();
let mut rng = StdRng::seed_from_u64(8765309);
let distribution = UniformMeshSampler::try_new(triangles).unwrap();
// 10000 random points from the surface of my_mesh:
let sample_points: Vec<Vec3> = distribution.sample_iter(&mut rng).take(10000).collect();
```

## Testing

Tested by instantiating meshes and sampling as demonstrated above.

---

## Changelog

- Added `Mesh::triangles` method to get a collection of triangles from a
mesh.
- Added `UniformMeshSampler` to `bevy_math::sampling`. This is a
distribution which allows random sampling over collections of triangles
(such as those provided through meshes).

---

## Discussion

### Design decisions

The main thing here was making sure to have a good separation between
the parts of this in `bevy_render` and in `bevy_math`. Getting the
triangles from a mesh seems like a reasonable step after adding
`Triangle3d` to `bevy_math`, so I decided to make all of the random
sampling operate at that level, with the fallible conversion to
triangles doing most of the work.

Notably, the sampler could be called something else that reflects that
its input is a collection of triangles, but if/when we add other kinds
of meshes to `bevy_math` (e.g. half-edge meshes), the fact that
`try_new` takes an `impl Into<Vec<Triangle3d>>` means that those meshes
just need to satisfy that trait bound in order to work immediately with
this sampling functionality. In that case, the result would just be
something like this:
```rust
let dist = UniformMeshSampler::try_new(mesh).unwrap();
```
I think this highlights that most of the friction is really just from
extracting data from `Mesh`.

It's maybe worth mentioning also that "collection of triangles"
(`Vec<Triangle3d>`) sits downstream of any other kind of triangle mesh,
since the topology connecting the triangles has been effectively erased,
which makes an `Into<Vec<Triangle3d>>` trait bound seem all the more
natural to me.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
This commit is contained in:
Matty 2024-07-07 20:57:08 -04:00 committed by GitHub
parent db19d2ee47
commit 900f50d77d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 191 additions and 2 deletions

View File

@ -18,6 +18,7 @@ approx = { version = "0.5", optional = true }
rand = { version = "0.8", features = [
"alloc",
], default-features = false, optional = true }
rand_distr = { version = "0.4.3", optional = true }
smallvec = { version = "1.11" }
bevy_reflect = { path = "../bevy_reflect", version = "0.14.0-dev", features = [
@ -49,7 +50,7 @@ glam_assert = ["glam/glam-assert"]
# Enable assertions in debug builds to check the validity of parameters passed to glam
debug_glam_assert = ["glam/debug-glam-assert"]
# Enable the rand dependency for shape_sampling
rand = ["dep:rand", "glam/rand"]
rand = ["dep:rand", "dep:rand_distr", "glam/rand"]
[lints]
workspace = true

View File

@ -0,0 +1,58 @@
//! Functionality related to random sampling from triangle meshes.
use crate::{
primitives::{Measured2d, Triangle3d},
ShapeSample, Vec3,
};
use rand::Rng;
use rand_distr::{Distribution, WeightedAliasIndex, WeightedError};
/// A [distribution] that caches data to allow fast sampling from a collection of triangles.
/// Generally used through [`sample`] or [`sample_iter`].
///
/// [distribution]: Distribution
/// [`sample`]: Distribution::sample
/// [`sample_iter`]: Distribution::sample_iter
///
/// Example
/// ```
/// # use bevy_math::{Vec3, primitives::*};
/// # use bevy_math::sampling::mesh_sampling::UniformMeshSampler;
/// # use rand::{SeedableRng, rngs::StdRng, distributions::Distribution};
/// let faces = Tetrahedron::default().faces();
/// let sampler = UniformMeshSampler::try_new(faces).unwrap();
/// let rng = StdRng::seed_from_u64(8765309);
/// // 50 random points on the tetrahedron:
/// let samples: Vec<Vec3> = sampler.sample_iter(rng).take(50).collect();
/// ```
pub struct UniformMeshSampler {
triangles: Vec<Triangle3d>,
face_distribution: WeightedAliasIndex<f32>,
}
impl Distribution<Vec3> for UniformMeshSampler {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Vec3 {
let face_index = self.face_distribution.sample(rng);
self.triangles[face_index].sample_interior(rng)
}
}
impl UniformMeshSampler {
/// Construct a new [`UniformMeshSampler`] from a list of [triangles].
///
/// Returns an error if the distribution of areas for the collection of triangles could not be formed
/// (most notably if the collection has zero surface area).
///
/// [triangles]: Triangle3d
pub fn try_new<T: IntoIterator<Item = Triangle3d>>(
triangles: T,
) -> Result<Self, WeightedError> {
let triangles: Vec<Triangle3d> = triangles.into_iter().collect();
let areas = triangles.iter().map(Measured2d::area).collect();
WeightedAliasIndex::new(areas).map(|face_distribution| Self {
triangles,
face_distribution,
})
}
}

View File

@ -2,8 +2,10 @@
//!
//! To use this, the "rand" feature must be enabled.
pub mod mesh_sampling;
pub mod shape_sampling;
pub mod standard;
pub use mesh_sampling::*;
pub use shape_sampling::*;
pub use standard::*;

View File

@ -18,7 +18,7 @@ use bevy_ecs::system::{
lifetimeless::{SRes, SResMut},
SystemParamItem,
};
use bevy_math::*;
use bevy_math::{primitives::Triangle3d, *};
use bevy_reflect::Reflect;
use bevy_utils::tracing::{error, warn};
use bytemuck::cast_slice;
@ -1076,6 +1076,134 @@ impl Mesh {
}
}
}
/// Get a list of this Mesh's [triangles] as an iterator if possible.
///
/// Returns an error if any of the following conditions are met (see [`MeshTrianglesError`]):
/// * The Mesh's [primitive topology] is not `TriangleList` or `TriangleStrip`.
/// * The Mesh is missing position or index data.
/// * The Mesh's position data has the wrong format (not `Float32x3`).
///
/// [primitive topology]: PrimitiveTopology
/// [triangles]: Triangle3d
pub fn triangles(&self) -> Result<impl Iterator<Item = Triangle3d> + '_, MeshTrianglesError> {
let Some(position_data) = self.attribute(Mesh::ATTRIBUTE_POSITION) else {
return Err(MeshTrianglesError::MissingPositions);
};
let Some(vertices) = position_data.as_float3() else {
return Err(MeshTrianglesError::PositionsFormat);
};
let Some(indices) = self.indices() else {
return Err(MeshTrianglesError::MissingIndices);
};
match self.primitive_topology {
PrimitiveTopology::TriangleList => {
// When indices reference out-of-bounds vertex data, the triangle is omitted.
// This implicitly truncates the indices to a multiple of 3.
let iterator = match indices {
Indices::U16(vec) => FourIterators::First(
vec.as_slice()
.chunks_exact(3)
.flat_map(move |indices| indices_to_triangle(vertices, indices)),
),
Indices::U32(vec) => FourIterators::Second(
vec.as_slice()
.chunks_exact(3)
.flat_map(move |indices| indices_to_triangle(vertices, indices)),
),
};
return Ok(iterator);
}
PrimitiveTopology::TriangleStrip => {
// When indices reference out-of-bounds vertex data, the triangle is omitted.
// If there aren't enough indices to make a triangle, then an empty vector will be
// returned.
let iterator = match indices {
Indices::U16(vec) => FourIterators::Third(
vec.as_slice()
.windows(3)
.flat_map(move |indices| indices_to_triangle(vertices, indices)),
),
Indices::U32(vec) => FourIterators::Fourth(
vec.as_slice()
.windows(3)
.flat_map(move |indices| indices_to_triangle(vertices, indices)),
),
};
return Ok(iterator);
}
_ => {
return Err(MeshTrianglesError::WrongTopology);
}
};
fn indices_to_triangle<T: TryInto<usize> + Copy>(
vertices: &[[f32; 3]],
indices: &[T],
) -> Option<Triangle3d> {
let vert0: Vec3 = Vec3::from(*vertices.get(indices[0].try_into().ok()?)?);
let vert1: Vec3 = Vec3::from(*vertices.get(indices[1].try_into().ok()?)?);
let vert2: Vec3 = Vec3::from(*vertices.get(indices[2].try_into().ok()?)?);
Some(Triangle3d {
vertices: [vert0, vert1, vert2],
})
}
}
}
/// A disjunction of four iterators. This is necessary to have a well-formed type for the output
/// of [`Mesh::triangles`], which produces iterators of four different types depending on the
/// branch taken.
enum FourIterators<A, B, C, D> {
First(A),
Second(B),
Third(C),
Fourth(D),
}
impl<A, B, C, D, I> Iterator for FourIterators<A, B, C, D>
where
A: Iterator<Item = I>,
B: Iterator<Item = I>,
C: Iterator<Item = I>,
D: Iterator<Item = I>,
{
type Item = I;
fn next(&mut self) -> Option<Self::Item> {
match self {
FourIterators::First(iter) => iter.next(),
FourIterators::Second(iter) => iter.next(),
FourIterators::Third(iter) => iter.next(),
FourIterators::Fourth(iter) => iter.next(),
}
}
}
/// An error that occurred while trying to extract a collection of triangles from a [`Mesh`].
#[derive(Debug, Error)]
pub enum MeshTrianglesError {
#[error("Source mesh does not have primitive topology TriangleList or TriangleStrip")]
WrongTopology,
#[error("Source mesh lacks position data")]
MissingPositions,
#[error("Source mesh position data is not Float32x3")]
PositionsFormat,
#[error("Source mesh lacks face index data")]
MissingIndices,
#[error("Face index data references vertices that do not exist")]
BadIndices,
}
impl core::ops::Mul<Mesh> for Transform {