From d2ef88f5e83a08159e18eae86fe214c4c4094454 Mon Sep 17 00:00:00 2001 From: Matty Date: Wed, 22 May 2024 08:38:08 -0400 Subject: [PATCH] Add `Distribution` access methods for `ShapeSample` trait (#13315) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stolen from #12835. # Objective Sometimes you want to sample a whole bunch of points from a shape instead of just one. You can write your own loop to do this, but it's really more idiomatic to use a `rand` [`Distribution`](https://docs.rs/rand/latest/rand/distributions/trait.Distribution.html) with the `sample_iter` method. Distributions also support other useful things like mapping, and they are suitable as generic items for consumption by other APIs. ## Solution `ShapeSample` has been given two new automatic trait methods, `interior_dist` and `boundary_dist`. They both have similar signatures (recall that `Output` is the output type for `ShapeSample`): ```rust fn interior_dist(self) -> impl Distribution where Self: Sized { //... } ``` These have default implementations which are powered by wrapper structs `InteriorOf` and `BoundaryOf` that actually implement `Distribution` — the implementations effectively just call `ShapeSample::sample_interior` and `ShapeSample::sample_boundary` on the contained type. The upshot is that this allows iteration as follows: ```rust // Get an iterator over boundary points of a rectangle: let rectangle = Rectangle::new(1.0, 2.0); let boundary_iter = rectangle.boundary_dist().sample_iter(rng); // Collect a bunch of boundary points at once: let boundary_pts: Vec = boundary_iter.take(1000).collect(); ``` Alternatively, you can use `InteriorOf`/`BoundaryOf` explicitly to similar effect: ```rust let boundary_pts: Vec = BoundaryOf(rectangle).sample_iter(rng).take(1000).collect(); ``` --- ## Changelog - Added `InteriorOf` and `BoundaryOf` distribution wrapper structs in `bevy_math::sampling::shape_sampling`. - Added `interior_dist` and `boundary_dist` automatic trait methods to `ShapeSample`. - Made `shape_sampling` module public with explanatory documentation. --- ## Discussion ### Design choices The main point of interest here is just the choice of `impl Distribution` instead of explicitly using `InteriorOf`/`BoundaryOf` return types for `interior_dist` and `boundary_dist`. The reason for this choice is that it allows future optimizations for repeated sampling — for example, instead of just wrapping the base type, `interior_dist`/`boundary_dist` could construct auxiliary data that is held over between sampling operations. --- crates/bevy_math/src/sampling/mod.rs | 2 +- .../bevy_math/src/sampling/shape_sampling.rs | 106 ++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) diff --git a/crates/bevy_math/src/sampling/mod.rs b/crates/bevy_math/src/sampling/mod.rs index 5944e5a355..3bb1b3fc51 100644 --- a/crates/bevy_math/src/sampling/mod.rs +++ b/crates/bevy_math/src/sampling/mod.rs @@ -2,7 +2,7 @@ //! //! To use this, the "rand" feature must be enabled. -mod shape_sampling; +pub mod shape_sampling; pub mod standard; pub use shape_sampling::*; diff --git a/crates/bevy_math/src/sampling/shape_sampling.rs b/crates/bevy_math/src/sampling/shape_sampling.rs index ed3359b832..520c1e87b7 100644 --- a/crates/bevy_math/src/sampling/shape_sampling.rs +++ b/crates/bevy_math/src/sampling/shape_sampling.rs @@ -1,3 +1,43 @@ +//! The [`ShapeSample`] trait, allowing random sampling from geometric shapes. +//! +//! At the most basic level, this allows sampling random points from the interior and boundary of +//! geometric primitives. For example: +//! ``` +//! # use bevy_math::primitives::*; +//! # use bevy_math::ShapeSample; +//! # use rand::SeedableRng; +//! # use rand::rngs::StdRng; +//! // Get some `Rng`: +//! let rng = &mut StdRng::from_entropy(); +//! // Make a circle of radius 2: +//! let circle = Circle::new(2.0); +//! // Get a point inside this circle uniformly at random: +//! let interior_pt = circle.sample_interior(rng); +//! // Get a point on the circle's boundary uniformly at random: +//! let boundary_pt = circle.sample_boundary(rng); +//! ``` +//! +//! For repeated sampling, `ShapeSample` also includes methods for accessing a [`Distribution`]: +//! ``` +//! # use bevy_math::primitives::*; +//! # use bevy_math::{Vec2, ShapeSample}; +//! # use rand::SeedableRng; +//! # use rand::rngs::StdRng; +//! # use rand::distributions::Distribution; +//! # let rng1 = StdRng::from_entropy(); +//! # let rng2 = StdRng::from_entropy(); +//! // Use a rectangle this time: +//! let rectangle = Rectangle::new(1.0, 2.0); +//! // Get an iterator that spits out random interior points: +//! let interior_iter = rectangle.interior_dist().sample_iter(rng1); +//! // Collect random interior points from the iterator: +//! let interior_pts: Vec = interior_iter.take(1000).collect(); +//! // Similarly, get an iterator over many random boundary points and collect them: +//! let boundary_pts: Vec = rectangle.boundary_dist().sample_iter(rng2).take(1000).collect(); +//! ``` +//! +//! In any case, the [`Rng`] used as the source of randomness must be provided explicitly. + use std::f32::consts::{PI, TAU}; use crate::{primitives::*, NormedVectorSpace, Vec2, Vec3}; @@ -39,6 +79,72 @@ pub trait ShapeSample { /// println!("{:?}", square.sample_boundary(&mut rand::thread_rng())); /// ``` fn sample_boundary(&self, rng: &mut R) -> Self::Output; + + /// Extract a [`Distribution`] whose samples are points of this shape's interior, taken uniformly. + /// + /// # Example + /// + /// ``` + /// # use bevy_math::prelude::*; + /// # use rand::distributions::Distribution; + /// let square = Rectangle::new(2.0, 2.0); + /// let rng = rand::thread_rng(); + /// + /// // Iterate over points randomly drawn from `square`'s interior: + /// for random_val in square.interior_dist().sample_iter(rng).take(5) { + /// println!("{:?}", random_val); + /// } + /// ``` + fn interior_dist(self) -> impl Distribution + where + Self: Sized, + { + InteriorOf(self) + } + + /// Extract a [`Distribution`] whose samples are points of this shape's boundary, taken uniformly. + /// + /// # Example + /// + /// ``` + /// # use bevy_math::prelude::*; + /// # use rand::distributions::Distribution; + /// let square = Rectangle::new(2.0, 2.0); + /// let rng = rand::thread_rng(); + /// + /// // Iterate over points randomly drawn from `square`'s boundary: + /// for random_val in square.boundary_dist().sample_iter(rng).take(5) { + /// println!("{:?}", random_val); + /// } + /// ``` + fn boundary_dist(self) -> impl Distribution + where + Self: Sized, + { + BoundaryOf(self) + } +} + +#[derive(Clone, Copy)] +/// A wrapper struct that allows interior sampling from a [`ShapeSample`] type directly as +/// a [`Distribution`]. +pub struct InteriorOf(pub T); + +#[derive(Clone, Copy)] +/// A wrapper struct that allows boundary sampling from a [`ShapeSample`] type directly as +/// a [`Distribution`]. +pub struct BoundaryOf(pub T); + +impl Distribution<::Output> for InteriorOf { + fn sample(&self, rng: &mut R) -> ::Output { + self.0.sample_interior(rng) + } +} + +impl Distribution<::Output> for BoundaryOf { + fn sample(&self, rng: &mut R) -> ::Output { + self.0.sample_boundary(rng) + } } impl ShapeSample for Circle {