hoomd_bevy/representation/
disk.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 circle.
10//!
11//! The [`Disk`] representation is a circle of pixels with a configurable
12//! outline color and an optional texture map.
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/disk.wgsl";
39
40/// Represent an entity with a 2D disk in the xy plane.
41///
42/// The base representation has a diameter of 1.0. Provide a non-unit diameter
43/// in [`sync`](Self::sync) to render disks of different sizes. Nominally, the z
44/// coordinate of the disks should be set to 0. Choose a different value to control
45/// 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 Disk<T> {
52    /// Mark the type of the disk.
53    marker: PhantomData<T>,
54}
55
56/// Assets that represent a Disk in the scene.
57#[derive(Resource)]
58pub struct Representation<T> {
59    /// The disk mesh.
60    mesh: Handle<Mesh>,
61    /// The disk material.
62    material: Handle<Material>,
63    /// Mark the type of the disk 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, "disk.wgsl");
79}
80
81impl<T: Send + Sync + 'static> Disk<T> {
82    /// Create assets to render disks.
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        asset_server: Res<AssetServer>,
92    ) {
93        #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
94        let background_colors = [material.0.background_color; 1024];
95
96        #[cfg(not(all(target_arch = "wasm32", not(feature = "webgpu"))))]
97        let background_colors =
98            buffers.add(ShaderStorageBuffer::from([material.0.background_color]));
99
100        let mesh = meshes.add(Rectangle::new(1.0, 1.0));
101        let material = Material {
102            background_colors,
103            #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
104            n_background_colors: 1,
105            outline_color: material.0.outline_color,
106            outline_width: material.0.outline_width,
107            texture_scale: material.0.texture_scale,
108            texture: material.0.texture_asset.map(|t| asset_server.load(t)),
109        };
110        let material = materials.add(material);
111
112        commands.insert_resource(Representation::<T> {
113            mesh,
114            material,
115            marker: PhantomData,
116        });
117    }
118
119    /// Copy the current positions of simulation particles to bevy entities.
120    pub fn sync<I>(
121        commands: &mut Commands,
122        disk_representation: Res<Representation<T>>,
123        query: Query<(Entity, &mut Transform), With<Self>>,
124        disks: I,
125    ) where
126        I: IntoIterator<Item = (Vec3, f32)>,
127    {
128        for (tag, item) in &mut query.into_iter().zip_longest(disks).enumerate() {
129            match item {
130                Both((_, mut transform), (position, diameter)) => {
131                    transform.translation = position;
132                    transform.scale = Vec3::splat(diameter);
133                }
134                Left((entity, _)) => commands.entity(entity).despawn(),
135                Right((position, diameter)) => {
136                    commands.spawn((
137                        MeshTag(tag as u32),
138                        Mesh2d(disk_representation.mesh.clone()),
139                        MeshMaterial2d(disk_representation.material.clone()),
140                        Transform::from_translation(position).with_scale(Vec3::splat(diameter)),
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 disk.
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    /// Factor to scale the texture by.
163    pub texture_scale: f32,
164
165    /// Name of the texture asset.
166    pub texture_asset: Option<String>,
167}
168
169impl Default for MaterialParameters {
170    fn default() -> Self {
171        Self {
172            background_color: PRIMARY_COLOR.into(),
173            outline_color: Color::linear_rgb(0.0, 0.0, 0.0).into(),
174            outline_width: 0.05,
175            texture_asset: None,
176            texture_scale: 1.2,
177        }
178    }
179}
180
181/// Control how disks are rendered.
182///
183/// Disks are always opaque and alpha in any texture or background color is ignored.
184///
185/// By default [`Material`] is initialized with only one background
186/// color. Color the instances differently by setting more than one color
187/// with [`set_background_colors`]. The color of each disk is given by
188/// `background_colors[tag % len(background_colors)]` so you may set fewer colors
189/// than there are disks. [`sync`] assigns `tag` values in increasing order to each
190/// primitive.
191///
192/// The `background_color` tints the texture by multiplication. With a `None`
193/// texture (the default), `background_color` sets the exact color of the disk.
194///
195/// Set the initial material by piping `MaterialParameters` into [`Disk::setup`].
196/// After it is initialized, change the material during execution via the `material`
197/// field in`ResMut<disk::Representation<A>>`.
198///
199/// [`sync`]: Disk::sync
200/// [`set_background_colors`]: Material::set_background_colors
201#[derive(Asset, TypePath, AsBindGroup, Debug, Clone)]
202pub struct Material {
203    /// Color applied to the outline.
204    #[uniform(0)]
205    outline_color: LinearRgba,
206
207    /// Width of the outline.
208    #[uniform(0)]
209    outline_width: f32,
210
211    /// Factor to scale the texture by.
212    #[uniform(0)]
213    texture_scale: f32,
214
215    /// Number of background colors in fixed size array.
216    #[uniform(0)]
217    #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
218    n_background_colors: u32,
219
220    /// Texture to apply. Tinted by `background_color`.
221    #[texture(1)]
222    #[sampler(2)]
223    texture: Option<Handle<Image>>,
224
225    /// Color applied to the interior of the disk (indexed by disk % array size).
226    #[uniform(3)]
227    #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
228    background_colors: [LinearRgba; 1024],
229
230    /// Color applied to the interior of the disk (indexed by disk % array size).
231    #[cfg(not(all(target_arch = "wasm32", not(feature = "webgpu"))))]
232    #[storage(3, read_only)]
233    background_colors: Handle<ShaderStorageBuffer>,
234}
235
236impl Material {
237    /// Set new background colors.
238    ///
239    /// # Panics
240    ///
241    /// WebGL2 builds (identified by the `wasm32` target without the `webgpu`
242    /// feature) support only 1024 background colors.
243    ///
244    /// Desktop target builds or `wasm32` target builds with `webgpu` support
245    /// a much larger number of colors and will not panic.
246    pub fn set_background_colors(
247        &mut self,
248        #[allow(
249            unused_variables,
250            unused_mut,
251            reason = "Not used in all build configurations."
252        )]
253        mut buffers: ResMut<Assets<ShaderStorageBuffer>>,
254        colors: &[LinearRgba],
255    ) {
256        #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
257        {
258            assert!(
259                colors.len() <= 1024,
260                "webgl2 builds support up to 1024 colors, got {}",
261                colors.len()
262            );
263            self.background_colors[..colors.len()].copy_from_slice(colors);
264            self.n_background_colors = colors.len() as u32;
265        }
266
267        #[cfg(not(all(target_arch = "wasm32", not(feature = "webgpu"))))]
268        {
269            let color_buffer = buffers
270                .get_mut(&self.background_colors)
271                .expect("Disk::setup should have added the storage buffer");
272
273            color_buffer.set_data(colors);
274        }
275    }
276}
277
278impl Material2d for Material {
279    fn fragment_shader() -> ShaderRef {
280        SHADER_ASSET_PATH.into()
281    }
282
283    fn vertex_shader() -> ShaderRef {
284        SHADER_ASSET_PATH.into()
285    }
286
287    fn alpha_mode(&self) -> AlphaMode2d {
288        AlphaMode2d::Mask(0.5)
289    }
290
291    #[cfg(all(target_arch = "wasm32", not(feature = "webgpu")))]
292    fn specialize(
293        descriptor: &mut RenderPipelineDescriptor,
294        _layout: &MeshVertexBufferLayoutRef,
295        _key: Material2dKey<Self>,
296    ) -> Result<(), SpecializedMeshPipelineError> {
297        descriptor.vertex.shader_defs.push("WEBGL2".into());
298
299        Ok(())
300    }
301}
hoomd_bevy/representation/disk.rs

hoomd_bevy/representation/
disk.rs
