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

hoomd_bevy/representation/
disk.rs
