hoomd_bevy/representation/
ellipse.rs

1// Copyright (c) 2024-2026 The Regents of the University of Michigan.
2// Part of hoomd-rs, released under the BSD 3-Clause License.
3
4#![allow(
5    clippy::missing_docs_in_private_items,
6    reason = "clippy reports a false positive errors in this file"
7)]
8
9//! An outlined ellipse.
10//!
11//! The [`Ellipse`] representation is a ellipse of pixels with a configurable
12//! outline color. Each ellipse can have a different aspect ratio.
13
14use bevy::{
15    asset::embedded_asset,
16    mesh::MeshTag,
17    prelude::*,
18    reflect::TypePath,
19    render::{render_resource::AsBindGroup, storage::ShaderStorageBuffer},
20    shader::ShaderRef,
21    sprite_render::{AlphaMode2d, Material2d, Material2dPlugin},
22};
23#[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
24use bevy::{
25    mesh::MeshVertexBufferLayoutRef,
26    render::render_resource::{RenderPipelineDescriptor, SpecializedMeshPipelineError},
27    sprite_render::Material2dKey,
28};
29use itertools::{
30    EitherOrBoth::{Both, Left, Right},
31    Itertools,
32};
33use std::marker::PhantomData;
34
35use crate::PRIMARY_COLOR;
36
37/// Location of the shader implementation
38const SHADER_ASSET_PATH: &str = "embedded://hoomd_bevy/representation/ellipse.wgsl";
39
40/// Represent an entity with a 2D ellipse in the xy plane.
41///
42/// The base representation has semi-axes (0.5, 0.5). Provide per-item axes
43/// in [`sync`](Self::sync) to render ellipses of different sizes and aspect ratios.
44/// Nominally, the z coordinate of the ellipses should be set to 0. Choose a different
45/// value to control the back to front draw order.
46///
47/// To use:
48/// * Add [`setup`](Self::setup) to the `Startup` schedule.
49/// * Call [`sync`](Self::sync) in an `Update` schedule that runs after `AdvanceSet`.
50#[derive(Component)]
51pub struct Ellipse<T> {
52    /// Mark the type of the ellipse.
53    marker: PhantomData<T>,
54}
55
56/// Assets that represent a ellipse in the scene.
57#[derive(Resource)]
58pub struct Representation<T> {
59    /// The ellipse mesh.
60    mesh: Handle<Mesh>,
61    /// The ellipse material.
62    material: Handle<Material>,
63    /// Mark the type of the ellipse assets.
64    marker: PhantomData<T>,
65}
66
67impl<T> Representation<T> {
68    /// Get the material
69    #[must_use]
70    pub fn material(&self) -> &Handle<Material> {
71        &self.material
72    }
73}
74
75/// Initialize needed plugins and add assets for this representation.
76pub(crate) fn build(app: &mut App) {
77    app.add_plugins(Material2dPlugin::<Material>::default());
78    embedded_asset!(app, "ellipse.wgsl");
79}
80
81impl<T: Send + Sync + 'static> Ellipse<T> {
82    /// Create assets to render ellipses.
83    pub fn setup(
84        material: In<MaterialParameters>,
85        mut commands: Commands,
86        #[cfg(not(all(target_arch = "wasm32", not(feature = "webgpu"))))] mut buffers: ResMut<
87            Assets<ShaderStorageBuffer>,
88        >,
89        mut meshes: ResMut<Assets<Mesh>>,
90        mut materials: ResMut<Assets<Material>>,
91    ) {
92        #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
93        let background_colors = [material.0.background_color; 1024];
94
95        #[cfg(not(all(target_arch = "wasm32", not(feature = "webgpu"))))]
96        let background_colors =
97            buffers.add(ShaderStorageBuffer::from([material.0.background_color]));
98
99        let mesh = meshes.add(Rectangle::new(1.0, 1.0));
100        let material = Material {
101            background_colors,
102            #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
103            n_background_colors: 1,
104            outline_color: material.0.outline_color,
105            outline_width: material.0.outline_width,
106        };
107        let material = materials.add(material);
108
109        commands.insert_resource(Representation::<T> {
110            mesh,
111            material,
112            marker: PhantomData,
113        });
114    }
115
116    /// Copy the current positions of simulation particles to bevy entities.
117    pub fn sync<I>(
118        commands: &mut Commands,
119        ellipse_representation: Res<Representation<T>>,
120        query: Query<(Entity, &mut Transform), With<Self>>,
121        ellipses: I,
122    ) where
123        I: IntoIterator<Item = (Vec3, f32, f32, f32)>,
124    {
125        for (tag, item) in &mut query.into_iter().zip_longest(ellipses).enumerate() {
126            match item {
127                Both((_, mut transform), (position, theta, a, b)) => {
128                    transform.translation = position;
129                    transform.rotation = Quat::from_rotation_z(theta);
130                    transform.scale = Vec3::new(a, b, 1.0);
131                }
132                Left((entity, _)) => commands.entity(entity).despawn(),
133                Right((position, theta, a, b)) => {
134                    commands.spawn((
135                        MeshTag(tag as u32),
136                        Mesh2d(ellipse_representation.mesh.clone()),
137                        MeshMaterial2d(ellipse_representation.material.clone()),
138                        Transform::from_translation(position)
139                            .with_scale(Vec3::new(a, b, 1.0))
140                            .with_rotation(Quat::from_rotation_z(theta)),
141                        Self {
142                            marker: PhantomData,
143                        },
144                    ));
145                }
146            }
147        }
148    }
149}
150
151/// Initialize [`Material`] with these settings.
152pub struct MaterialParameters {
153    /// Color applied to the interior of the ellipses.
154    pub background_color: LinearRgba,
155
156    /// Color applied to the outline.
157    pub outline_color: LinearRgba,
158
159    /// Width of the outline.
160    pub outline_width: f32,
161}
162
163impl Default for MaterialParameters {
164    fn default() -> Self {
165        Self {
166            background_color: PRIMARY_COLOR.into(),
167            outline_color: Color::linear_rgb(0.0, 0.0, 0.0).into(),
168            outline_width: 0.05,
169        }
170    }
171}
172
173/// Control how ellipses are rendered.
174///
175/// Ellipses are always opaque and alpha in any background color is ignored.
176///
177/// By default [`Material`] is initialized with only one background
178/// color. Color the instances differently by setting more than one color
179/// with [`set_background_colors`]. The color of each ellipse is given by
180/// `background_colors[tag % len(background_colors)]` so you may set fewer colors
181/// than there are ellipses. [`sync`] assigns `tag` values in increasing order to each
182/// primitive.
183///
184/// The `background_color` tints the texture by multiplication. With a `None`
185/// texture (the default), `background_color` sets the exact color of the ellipses.
186///
187/// Set the initial material by piping `MaterialParameters` into [`Ellipse::setup`].
188/// After it is initialized, change the material during execution via the `material`
189/// field in`ResMut<ellipse::Representation<A>>`.
190///
191/// [`sync`]: Ellipse::sync
192/// [`set_background_colors`]: Material::set_background_colors
193#[derive(Asset, TypePath, AsBindGroup, Debug, Clone)]
194pub struct Material {
195    /// Color applied to the outline.
196    #[uniform(0)]
197    outline_color: LinearRgba,
198
199    /// Width of the outline.
200    #[uniform(0)]
201    outline_width: f32,
202
203    /// Number of background colors in fixed size array.
204    #[uniform(0)]
205    #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
206    n_background_colors: u32,
207
208    /// Color applied to the interior of the ellipse (indexed by ellipse % array size).
209    #[uniform(1)]
210    #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
211    background_colors: [LinearRgba; 1024],
212
213    /// Color applied to the interior of the ellipse (indexed by ellipse % array size).
214    #[cfg(not(all(target_arch = "wasm32", not(feature = "webgpu"))))]
215    #[storage(1, read_only)]
216    background_colors: Handle<ShaderStorageBuffer>,
217}
218
219impl Material {
220    /// Set new background colors.
221    ///
222    /// # Panics
223    ///
224    /// WebGL2 builds (identified by the `wasm32` target without the `webgpu`
225    /// feature) support only 1024 background colors.
226    ///
227    /// Desktop target builds or `wasm32` target builds with `webgpu` support
228    /// a much larger number of colors and will not panic.
229    pub fn set_background_colors(
230        &mut self,
231        #[allow(
232            unused_variables,
233            unused_mut,
234            reason = "Not used in all build configurations."
235        )]
236        mut buffers: ResMut<Assets<ShaderStorageBuffer>>,
237        colors: &[LinearRgba],
238    ) {
239        #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
240        {
241            assert!(
242                colors.len() <= 1024,
243                "webgl2 builds support up to 1024 colors, got {}",
244                colors.len()
245            );
246            self.background_colors[..colors.len()].copy_from_slice(colors);
247            self.n_background_colors = colors.len() as u32;
248        }
249
250        #[cfg(not(all(target_arch = "wasm32", not(feature = "webgpu"))))]
251        {
252            let color_buffer = buffers
253                .get_mut(&self.background_colors)
254                .expect("Ellipse::setup should have added the storage buffer");
255
256            color_buffer.set_data(colors);
257        }
258    }
259}
260
261impl Material2d for Material {
262    fn fragment_shader() -> ShaderRef {
263        SHADER_ASSET_PATH.into()
264    }
265
266    fn vertex_shader() -> ShaderRef {
267        SHADER_ASSET_PATH.into()
268    }
269
270    fn alpha_mode(&self) -> AlphaMode2d {
271        AlphaMode2d::Mask(0.5)
272    }
273
274    #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
275    fn specialize(
276        descriptor: &mut RenderPipelineDescriptor,
277        _layout: &MeshVertexBufferLayoutRef,
278        _key: Material2dKey<Self>,
279    ) -> Result<(), SpecializedMeshPipelineError> {
280        descriptor.vertex.shader_defs.push("WEBGL2".into());
281
282        Ok(())
283    }
284}
hoomd_bevy/representation/ellipse.rs

hoomd_bevy/representation/
ellipse.rs
