Skip to content

Commit

Permalink
Add force and its example
Browse files Browse the repository at this point in the history
  • Loading branch information
narasan49 committed Aug 14, 2024
1 parent 5d4bb67 commit f7132cf
Show file tree
Hide file tree
Showing 4 changed files with 289 additions and 7 deletions.
41 changes: 41 additions & 0 deletions assets/shaders/add_force.wgsl
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
#import bevy_fluid::fluid_uniform::SimulationUniform;

@group(0) @binding(0) var<storage, read> force: array<vec2<f32>>;
@group(0) @binding(1) var<storage, read> position: array<vec2<f32>>;
@group(0) @binding(2) var u: texture_storage_2d<r32float, read_write>;
@group(0) @binding(3) var v: texture_storage_2d<r32float, read_write>;

@group(1) @binding(0) var<uniform> constants: SimulationUniform;

@compute
@workgroup_size(1, 64, 1)
fn add_force(
@builtin(global_invocation_id) invocation_id: vec3<u32>,
) {
let x_u = vec2<i32>(i32(invocation_id.x), i32(invocation_id.y));
let x_v = vec2<i32>(x_u.y, x_u.x);
var n = arrayLength(&force);
var net_force = vec2<f32>(0.0, 0.0);

loop {
if (n == 0) {
break;
}
n = n - 1u;
let f = force[n];
let p = position[n];
let force_u = f.x * gaussian_2d(vec2<f32>(x_u), p, 10.0);
let force_v = f.y * gaussian_2d(vec2<f32>(x_v), p, 10.0);
net_force = net_force + vec2<f32>(force_u, force_v);
}

let u_val = textureLoad(u, x_u).r;
let v_val = textureLoad(v, x_v).r;
textureStore(u, x_u, vec4<f32>(u_val + net_force.x * constants.dt, 0.0, 0.0, 0.0));
textureStore(v, x_v, vec4<f32>(v_val + net_force.y * constants.dt, 0.0, 0.0, 0.0));
}

fn gaussian_2d(x: vec2<f32>, x0: vec2<f32>, sigma: f32) -> f32 {
let b = -1.0 / (2.0 * sigma * sigma);
return exp(b * dot(x - x0, x - x0));
}
120 changes: 120 additions & 0 deletions examples/mouse_interaction.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,120 @@
extern crate bevy_fluid;

// use advection_plugin::AdvectionPlugin;
use bevy::{
asset::AssetMetaCheck,
input::mouse::MouseMotion,
prelude::*,
render::{
settings::{Backends, WgpuSettings},
RenderPlugin,
},
sprite::MaterialMesh2dBundle,
window::PrimaryWindow,
};

use bevy_fluid::euler_fluid::{
add_force::AddForceMaterial, advection::AdvectionMaterial, fluid_material::VelocityMaterial,
FluidPlugin,
};

const WIDTH: f32 = 1280.0;
const HEIGHT: f32 = 720.0;

fn main() {
let mut app = App::new();
// [workaround] Asset meta files cannot be found on browser.
// see also: https://github.com/bevyengine/bevy/issues/10157
let meta_check = if cfg!(target_arch = "wasm32") {
AssetMetaCheck::Never
} else {
AssetMetaCheck::Always
};

app.add_plugins(
DefaultPlugins
.set(WindowPlugin {
primary_window: Some(Window {
resolution: (WIDTH, HEIGHT).into(),
title: "bevy fluid".to_string(),
..default()
}),
..default()
})
.set(RenderPlugin {
render_creation: bevy::render::settings::RenderCreation::Automatic(WgpuSettings {
backends: Some(Backends::DX12 | Backends::BROWSER_WEBGPU),
..default()
}),
..default()
})
.set(AssetPlugin {
meta_check,
..default()
}),
)
.add_plugins(FluidPlugin)
.add_systems(Startup, setup_scene)
.add_systems(Update, on_advection_initialized)
.add_systems(Update, mouse_motion);

app.run();
}

fn setup_scene(mut commands: Commands) {
commands
.spawn(Camera2dBundle::default())
.insert(Name::new("Camera"));
}

fn on_advection_initialized(
mut commands: Commands,
advection: Option<Res<AdvectionMaterial>>,
mut meshes: ResMut<Assets<Mesh>>,
mut materials: ResMut<Assets<VelocityMaterial>>,
) {
if let Some(advection) = advection {
if advection.is_changed() {
info!("prepare velocity texture");
// spwan plane to visualize advection
let mesh = meshes.add(Rectangle::default());
let material = materials.add(VelocityMaterial {
offset: 0.5,
scale: 0.1,
u: Some(advection.u_in.clone()),
v: Some(advection.v_in.clone()),
});

commands.spawn(MaterialMesh2dBundle {
mesh: mesh.into(),
transform: Transform::default().with_scale(Vec3::splat(512.)),
material,
..default()
});
}
}
}

fn mouse_motion(
mouse_button_input: Res<ButtonInput<MouseButton>>,
mut mouse_motion: EventReader<MouseMotion>,
mut force_material: ResMut<AddForceMaterial>,
q_window: Query<&Window, With<PrimaryWindow>>,
) {
if mouse_button_input.pressed(MouseButton::Left) {
if let Some(cursor_position) = q_window.single().cursor_position() {
let force = mouse_motion
.read()
.map(|mouse| mouse.delta)
.collect::<Vec<_>>();
let position = vec![cursor_position; force.len()];
force_material.force = force;
force_material.position = position;

return;
}
}

force_material.force = vec![];
force_material.position = vec![];
}
43 changes: 36 additions & 7 deletions src/euler_fluid.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
pub mod add_force;
pub mod advection;
pub mod fluid_material;
pub mod geometry;
pub mod grid_label;
pub mod projection;
pub mod uniform;

use add_force::{AddForceBindGroup, AddForceMaterial, AddForcePipeline};
use advection::{AdvectionBindGroup, AdvectionMaterial, AdvectionPipeline};
use bevy::{
asset::load_internal_asset,
Expand All @@ -19,7 +21,9 @@ use bevy::{
renderer::RenderDevice,
Render, RenderApp, RenderSet,
},
sprite::Material2dPlugin,
};
use fluid_material::VelocityMaterial;
use geometry::{CircleCollectionBindGroup, CircleCollectionMaterial, CrircleUniform, Velocity};
use grid_label::{GridLabelBindGroup, GridLabelMaterial, GridLabelPipeline};
use projection::{
Expand All @@ -41,15 +45,15 @@ const WORKGROUP_SIZE: u32 = 8;
const FLUID_UNIFORM_SHADER_HANDLE: Handle<Shader> =
Handle::weak_from_u128(0x8B9323522322463BA8CF530771C532EF);

const COORDINATE_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(0x9F8E2E5B1E5F40C096C31175C285BF11);
const COORDINATE_SHADER_HANDLE: Handle<Shader> =
Handle::weak_from_u128(0x9F8E2E5B1E5F40C096C31175C285BF11);

pub struct FluidPlugin;

impl Plugin for FluidPlugin {
fn build(&self, app: &mut App) {
app
// .add_plugins(ExtractResourcePlugin::<VelocityTexture>::default())
.add_plugins(ExtractResourcePlugin::<AdvectionMaterial>::default())
app.add_plugins(ExtractResourcePlugin::<AdvectionMaterial>::default())
.add_plugins(ExtractResourcePlugin::<AddForceMaterial>::default())
.add_plugins(ExtractResourcePlugin::<DivergenceMaterial>::default())
.add_plugins(ExtractResourcePlugin::<SolvePressureMaterial>::default())
.add_plugins(ExtractResourcePlugin::<JacobiMaterial>::default())
Expand All @@ -69,6 +73,10 @@ impl Plugin for FluidPlugin {
Render,
prepare_bind_group.in_set(RenderSet::PrepareBindGroups),
)
.add_systems(
Render,
add_force::prepare_bind_group.in_set(RenderSet::PrepareBindGroups),
)
.add_systems(
Render,
advection::prepare_bind_group.in_set(RenderSet::PrepareBindGroups),
Expand Down Expand Up @@ -116,7 +124,7 @@ impl Plugin for FluidPlugin {
fn finish(&self, app: &mut App) {
let render_app = app.sub_app_mut(RenderApp);
render_app.init_resource::<AdvectionPipeline>();
// render_app.init_resource::<FluidPipeline>();
render_app.init_resource::<AddForcePipeline>();
render_app.init_resource::<SolvePressurePipeline>();
render_app.init_resource::<DivergencePipeline>();
render_app.init_resource::<JacobiPipeline>();
Expand Down Expand Up @@ -173,6 +181,12 @@ fn setup(mut commands: Commands, mut images: ResMut<Assets<Image>>) {
v_in: v0.clone(),
v_out: v1.clone(),
});
commands.insert_resource(AddForceMaterial {
force: vec![],
position: vec![],
u: u1.clone(),
v: v1.clone(),
});
commands.insert_resource(DivergenceMaterial {
u: u1.clone(),
v: v1.clone(),
Expand All @@ -199,14 +213,14 @@ fn setup(mut commands: Commands, mut images: ResMut<Assets<Image>>) {

commands.spawn(SimulationUniform {
dx: 1.0f32,
dt: 0.1f32,
dt: 0.5f32,
rho: 1.293f32,
});
}

fn update(mut query: Query<&mut SimulationUniform>, _time: Res<Time>) {
for mut uniform in &mut query {
uniform.dt = 0.1;
uniform.dt = 0.5;
}
}

Expand Down Expand Up @@ -253,6 +267,7 @@ impl render_graph::Node for FluidNode {
// update node state
fn update(&mut self, world: &mut World) {
let advection_pipeline = world.resource::<AdvectionPipeline>();
let add_force_pipeline = world.resource::<AddForcePipeline>();
let divergence_pipeline = world.resource::<DivergencePipeline>();
let jacobi_pipeline = world.resource::<JacobiPipeline>();
let solve_pipeline = world.resource::<SolvePressurePipeline>();
Expand All @@ -275,6 +290,8 @@ impl render_graph::Node for FluidNode {
FluidState::Init => {
let advection_pipeline =
pipeline_cache.get_compute_pipeline_state(advection_pipeline.pipeline);
let add_force_pipeline =
pipeline_cache.get_compute_pipeline_state(add_force_pipeline.pipeline);
let jacobi_pipeline_state =
pipeline_cache.get_compute_pipeline_state(jacobi_pipeline.pipeline);
let jacobi_swap_pipeline =
Expand All @@ -287,6 +304,7 @@ impl render_graph::Node for FluidNode {
pipeline_cache.get_compute_pipeline_state(grid_label_pipeline.update_pipeline);
match (
advection_pipeline,
add_force_pipeline,
jacobi_pipeline_state,
jacobi_swap_pipeline,
solve_pipeline,
Expand All @@ -300,6 +318,7 @@ impl render_graph::Node for FluidNode {
CachedPipelineState::Ok(_),
CachedPipelineState::Ok(_),
CachedPipelineState::Ok(_),
CachedPipelineState::Ok(_),
) => self.state = FluidState::Update,
_ => {}
}
Expand Down Expand Up @@ -366,6 +385,16 @@ impl render_graph::Node for FluidNode {
pass.set_bind_group(2, grid_label_bind_group, &[]);
pass.dispatch_workgroups(SIZE.0 + 1, SIZE.1 / WORKGROUP_SIZE / WORKGROUP_SIZE, 1);

// Add force if triggered.
let add_force_pipeline = world.resource::<AddForcePipeline>();
let add_force_pipeline = pipeline_cache
.get_compute_pipeline(add_force_pipeline.pipeline)
.unwrap();
let add_force_bind_group = &world.resource::<AddForceBindGroup>().0;
pass.set_pipeline(add_force_pipeline);
pass.set_bind_group(0, add_force_bind_group, &vec![]);
pass.dispatch_workgroups(SIZE.0 + 1, SIZE.1 / WORKGROUP_SIZE / WORKGROUP_SIZE, 1);

let divergence_pipeline = world.resource::<DivergencePipeline>();
let divergence_pipeline = pipeline_cache
.get_compute_pipeline(divergence_pipeline.pipeline)
Expand Down
92 changes: 92 additions & 0 deletions src/euler_fluid/add_force.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
use std::borrow::Cow;

use bevy::{
prelude::*,
render::{
extract_resource::ExtractResource,
render_asset::RenderAssets,
render_resource::{
binding_types::uniform_buffer, AsBindGroup, BindGroup, BindGroupLayout,
BindGroupLayoutEntries, CachedComputePipelineId, ComputePipelineDescriptor,
PipelineCache, ShaderStages,
},
renderer::RenderDevice,
texture::{FallbackImage, GpuImage},
},
};

use super::uniform::SimulationUniform;

pub fn prepare_bind_group(
mut commands: Commands,
gpu_images: Res<RenderAssets<GpuImage>>,
pipeline: Res<AddForcePipeline>,
textures: Res<AddForceMaterial>,
render_device: Res<RenderDevice>,
fallback_image: Res<FallbackImage>,
) {
let bind_group = textures
.as_bind_group(
&pipeline.bind_group_layout,
&render_device,
&gpu_images,
&fallback_image,
)
.unwrap()
.bind_group;

commands.insert_resource(AddForceBindGroup(bind_group));
}

#[derive(Resource, Clone, ExtractResource, AsBindGroup)]
pub struct AddForceMaterial {
#[storage(0, read_only, visibility(compute))]
pub force: Vec<Vec2>,
#[storage(1, read_only, visibility(compute))]
pub position: Vec<Vec2>,
#[storage_texture(2, image_format = R32Float, access = ReadWrite)]
pub u: Handle<Image>,
#[storage_texture(3, image_format = R32Float, access = ReadWrite)]
pub v: Handle<Image>,
}

#[derive(Resource)]
pub struct AddForceBindGroup(pub BindGroup);

#[derive(Resource)]
pub struct AddForcePipeline {
bind_group_layout: BindGroupLayout,
pub pipeline: CachedComputePipelineId,
}

impl FromWorld for AddForcePipeline {
fn from_world(world: &mut World) -> Self {
let render_device = world.resource::<RenderDevice>();
let bind_group_layout = AddForceMaterial::bind_group_layout(render_device);
let uniform_bind_group_layout = render_device.create_bind_group_layout(
None,
&BindGroupLayoutEntries::single(
ShaderStages::COMPUTE,
uniform_buffer::<SimulationUniform>(false),
),
);

let shader = world
.resource::<AssetServer>()
.load("shaders/add_force.wgsl");
let pipeline_cache = world.resource::<PipelineCache>();
let pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
label: None,
layout: vec![bind_group_layout.clone(), uniform_bind_group_layout.clone()],
push_constant_ranges: Vec::new(),
shader: shader.clone(),
shader_defs: vec![],
entry_point: Cow::from("add_force"),
});

Self {
bind_group_layout,
pipeline,
}
}
}

0 comments on commit f7132cf

Please sign in to comment.