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::ShaderBuffer},
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<ShaderBuffer>,
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 = buffers.add(ShaderBuffer::from([material.0.background_color]));
97
98        let mesh = meshes.add(Rectangle::new(1.0, 1.0));
99        let material = Material {
100            background_colors,
101            #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
102            n_background_colors: 1,
103            outline_color: material.0.outline_color,
104            outline_width: material.0.outline_width,
105        };
106        let material = materials.add(material);
107
108        commands.insert_resource(Representation::<T> {
109            mesh,
110            material,
111            marker: PhantomData,
112        });
113    }
114
115    /// Copy the current positions of simulation particles to bevy entities.
116    pub fn sync<I>(
117        commands: &mut Commands,
118        ellipse_representation: Res<Representation<T>>,
119        query: Query<(Entity, &mut Transform), With<Self>>,
120        ellipses: I,
121    ) where
122        I: IntoIterator<Item = (Vec3, f32, f32, f32)>,
123    {
124        for (tag, item) in &mut query.into_iter().zip_longest(ellipses).enumerate() {
125            match item {
126                Both((_, mut transform), (position, theta, a, b)) => {
127                    transform.translation = position;
128                    transform.rotation = Quat::from_rotation_z(theta);
129                    transform.scale = Vec3::new(a, b, 1.0);
130                }
131                Left((entity, _)) => commands.entity(entity).despawn(),
132                Right((position, theta, a, b)) => {
133                    commands.spawn((
134                        MeshTag(tag as u32),
135                        Mesh2d(ellipse_representation.mesh.clone()),
136                        MeshMaterial2d(ellipse_representation.material.clone()),
137                        Transform::from_translation(position)
138                            .with_scale(Vec3::new(a, b, 1.0))
139                            .with_rotation(Quat::from_rotation_z(theta)),
140                        Self {
141                            marker: PhantomData,
142                        },
143                    ));
144                }
145            }
146        }
147    }
148}
149
150/// Initialize [`Material`] with these settings.
151pub struct MaterialParameters {
152    /// Color applied to the interior of the ellipses.
153    pub background_color: LinearRgba,
154
155    /// Color applied to the outline.
156    pub outline_color: LinearRgba,
157
158    /// Width of the outline.
159    pub outline_width: f32,
160}
161
162impl Default for MaterialParameters {
163    fn default() -> Self {
164        Self {
165            background_color: PRIMARY_COLOR.into(),
166            outline_color: Color::linear_rgb(0.0, 0.0, 0.0).into(),
167            outline_width: 0.05,
168        }
169    }
170}
171
172/// Control how ellipses are rendered.
173///
174/// Ellipses are always opaque and alpha in any background color is ignored.
175///
176/// By default [`Material`] is initialized with only one background
177/// color. Color the instances differently by setting more than one color
178/// with [`set_background_colors`]. The color of each ellipse is given by
179/// `background_colors[tag % len(background_colors)]` so you may set fewer colors
180/// than there are ellipses. [`sync`] assigns `tag` values in increasing order to each
181/// primitive.
182///
183/// The `background_color` tints the texture by multiplication. With a `None`
184/// texture (the default), `background_color` sets the exact color of the ellipses.
185///
186/// Set the initial material by piping `MaterialParameters` into [`Ellipse::setup`].
187/// After it is initialized, change the material during execution via the `material`
188/// field in`ResMut<ellipse::Representation<A>>`.
189///
190/// [`sync`]: Ellipse::sync
191/// [`set_background_colors`]: Material::set_background_colors
192#[derive(Asset, TypePath, AsBindGroup, Debug, Clone)]
193pub struct Material {
194    /// Color applied to the outline.
195    #[uniform(0)]
196    outline_color: LinearRgba,
197
198    /// Width of the outline.
199    #[uniform(0)]
200    outline_width: f32,
201
202    /// Number of background colors in fixed size array.
203    #[uniform(0)]
204    #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
205    n_background_colors: u32,
206
207    /// Color applied to the interior of the ellipse (indexed by ellipse % array size).
208    #[uniform(1)]
209    #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
210    background_colors: [LinearRgba; 1024],
211
212    /// Color applied to the interior of the ellipse (indexed by ellipse % array size).
213    #[cfg(not(all(target_arch = "wasm32", not(feature = "webgpu"))))]
214    #[storage(1, read_only)]
215    background_colors: Handle<ShaderBuffer>,
216}
217
218impl Material {
219    /// Set new background colors.
220    ///
221    /// # Panics
222    ///
223    /// WebGL2 builds (identified by the `wasm32` target without the `webgpu`
224    /// feature) support only 1024 background colors.
225    ///
226    /// Desktop target builds or `wasm32` target builds with `webgpu` support
227    /// a much larger number of colors and will not panic.
228    pub fn set_background_colors(
229        &mut self,
230        #[allow(
231            unused_variables,
232            unused_mut,
233            reason = "Not used in all build configurations."
234        )]
235        mut buffers: ResMut<Assets<ShaderBuffer>>,
236        colors: &[LinearRgba],
237    ) {
238        #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
239        {
240            assert!(
241                colors.len() <= 1024,
242                "webgl2 builds support up to 1024 colors, got {}",
243                colors.len()
244            );
245            self.background_colors[..colors.len()].copy_from_slice(colors);
246            self.n_background_colors = colors.len() as u32;
247        }
248
249        #[cfg(not(all(target_arch = "wasm32", not(feature = "webgpu"))))]
250        {
251            let mut color_buffer = buffers
252                .get_mut(&self.background_colors)
253                .expect("Ellipse::setup should have added the storage buffer");
254
255            color_buffer.set_data(colors);
256        }
257    }
258}
259
260impl Material2d for Material {
261    fn fragment_shader() -> ShaderRef {
262        SHADER_ASSET_PATH.into()
263    }
264
265    fn vertex_shader() -> ShaderRef {
266        SHADER_ASSET_PATH.into()
267    }
268
269    fn alpha_mode(&self) -> AlphaMode2d {
270        AlphaMode2d::Mask(0.5)
271    }
272
273    #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
274    fn specialize(
275        descriptor: &mut RenderPipelineDescriptor,
276        _layout: &MeshVertexBufferLayoutRef,
277        _key: Material2dKey<Self>,
278    ) -> Result<(), SpecializedMeshPipelineError> {
279        descriptor.vertex.shader_defs.push("WEBGL2".into());
280
281        Ok(())
282    }
283}
hoomd_bevy/representation/ellipse.rs

hoomd_bevy/representation/
ellipse.rs
