Switch to using wgpu for vis

Cargo.toml

@@ -6,3 +6,9 @@ edition = "2024"
 [dependencies]
 crossterm = "0.29.0"
 ratatui = "0.29.0"
+
+env_logger = "0.11"
+pollster = "0.4"
+wgpu = "28.0.0"
+winit = { version = "0.30.8", features = ["android-native-activity"] }
+bytemuck = { version = "1.24", features = ["derive"] }

README.md

@@ -4,3 +4,10 @@ A simple 2D fluid simulation
 ## References
 - [Website from Ten Minute Physics](https://matthias-research.github.io/pages/tenMinutePhysics/17-fluidSim.html)
 - [PDF from Ten Minute Physics](https://matthias-research.github.io/pages/tenMinutePhysics/17-fluidSim.pdf)
+
+## Debug on Macos
+- Build in release mode
+- Start XCode
+- Debug Executable -> Select newly build binary
+- Start task
+- Click on Metal Logo and Capture a single frame

src/main.rs

@@ -1,96 +1,579 @@
-use crossterm::event::{self, Event};
-use ratatui::{text::Text, Frame, widgets::canvas::Canvas, widgets::canvas::Points, widgets::canvas::Rectangle};
+use std::sync::Arc;
+use std::borrow::Cow;
 
-use fluidsim::{Domain, simulate};
+use winit::{
+    application::ApplicationHandler,
+    event::WindowEvent,
+    keyboard::Key,
+    event_loop::{ActiveEventLoop, ControlFlow, EventLoop, OwnedDisplayHandle},
+    window::{Window, WindowId},
+};
+use wgpu::util::DeviceExt;
 
-fn main() {
-    let width = 140;
-    let height = 60;
+struct State {
+    window: Arc<Window>,
+    device: wgpu::Device,
+    queue: Arc<wgpu::Queue>,
+    size: winit::dpi::PhysicalSize<u32>,
+    surface: wgpu::Surface<'static>,
+    surface_format: wgpu::TextureFormat,
 
-    let mut t_domain = fluidsim::ten_minute_physics::Domain::new(1000.0, width, height);
-    t_domain.set_s_with(|x, y| {
-        1.0
-    });
-    t_domain.set_u_with(|x, y| {
-        if x == 0 {
-            2.0
-        } else {
-            0.0
-        }
-    });
+    render_pipeline: wgpu::RenderPipeline,
 
-    t_domain.set_s_with(|x, y| {
-        if x > 40 && x < 50 && y > 20 && y < 30 {
-            return 0.0;
-        }
-
-        1.0
-    });
-
-    let domain_mtx_1 = std::sync::Arc::new(std::sync::Mutex::new(t_domain.clone()));
-    let domain_mtx = domain_mtx_1.clone();
-    std::thread::spawn(move || {
-        let mut start = t_domain;
-
-        for i in 0.. {
-            // println!("Iteration {}", i);
-
-            start.reset_pressure();
-
-            start.solve_incompressibility(100, 0.01);
-            start.extrapolate();
-            start.advect_vel(0.01);
-            start.advect_smoke(0.01);
-
-            if let Ok(mut mtx) = domain_mtx.try_lock() {
-                *mtx = start.clone();
-            }
-
-            std::thread::sleep(std::time::Duration::from_millis(10));
-        }
-    });
-
-    let domain_mtx = domain_mtx_1;
-    let mut domain = domain_mtx.lock().unwrap().clone();
-
-    let mut terminal = ratatui::init();
-    loop {
-        domain = domain_mtx.lock().unwrap().clone();
-
-        terminal.draw(|frame| {
-            draw(frame, &domain)
-        }).expect("failed to draw frame");
-        match event::read().expect("failed to read event") {
-            Event::Key(k) if k.code == crossterm::event::KeyCode::Esc => {
-                break;
-            }
-            _ => {}
-        };
-    }
-    ratatui::restore();
+    vertices: Vec<Vertex>,
+    vertex_buffer: Arc<wgpu::Buffer>,
+    texture: Arc<wgpu::Texture>,
+    diffuse_bind_group: wgpu::BindGroup,
 }
 
-fn draw(frame: &mut Frame, domain: &fluidsim::ten_minute_physics::Domain) {
-    let (width, height) = domain.inner_dims();
+#[repr(C)]
+#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
+struct Vertex {
+    position: [f32; 3],
+    world_pos: [f32; 3],
+}
 
-    let canvas = Canvas::default().x_bounds([0.0, width as f64]).y_bounds([0.0, height as f64]).paint(|ctx| {
-        let (min_v, max_v) = domain.smoke_iter().fold((f32::MAX, f32::MIN), |(mip, map), v| {
-            (mip.min(v.2), map.max(v.2))
+impl Vertex {
+    fn desc() -> wgpu::VertexBufferLayout<'static> {
+        wgpu::VertexBufferLayout {
+            array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
+            step_mode: wgpu::VertexStepMode::Vertex,
+            attributes: &[
+                wgpu::VertexAttribute {
+                    offset: 0,
+                    shader_location: 0,
+                    format: wgpu::VertexFormat::Float32x3,
+                },
+                wgpu::VertexAttribute {
+                    offset: std::mem::size_of::<[f32; 3]>() as wgpu::BufferAddress,
+                    shader_location: 1,
+                    format: wgpu::VertexFormat::Float32x3,
+                }
+            ]
+        }
+    }
+}
+
+
+impl State {
+    async fn new(display: OwnedDisplayHandle, window: Arc<Window>, domain: &fluidsim::ten_minute_physics::Domain) -> State {
+        let instance = wgpu::Instance::new(
+            &wgpu::InstanceDescriptor::default(),
+        );
+        let adapter = instance
+            .request_adapter(&wgpu::RequestAdapterOptions::default())
+            .await
+            .unwrap();
+        let (device, queue) = adapter
+            .request_device(&wgpu::DeviceDescriptor::default())
+            .await
+            .unwrap();
+
+        let size = window.inner_size();
+
+        let surface = instance.create_surface(window.clone()).unwrap();
+        let cap = surface.get_capabilities(&adapter);
+        let surface_format = cap.formats[0];
+
+        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
+            label: Some("Shader"),
+            source: wgpu::ShaderSource::Wgsl(include_str!("shader.wgsl").into()),
+        });
+
+        let (ig_width, ig_height) = domain.inner_dims();
+        let (g_width, g_height) = (ig_width as f32+2.0, ig_height as f32 + 2.0);
+        let vertices: Vec<Vertex> = domain.all_grid_iter().flat_map(|(x, y)| {
+            let left_bound = ((x as f32 / g_width) - 0.5)*2.0;
+            let upper_bound = ((y as f32 / g_height) - 0.5)*-2.0;
+
+            let h_step = 2.0 / g_height;
+            let w_step = 2.0 / g_width;
+
+            [
+                Vertex {
+                    position: [left_bound, upper_bound, 0.0],
+                    world_pos: [x as f32 / g_width, y as f32 / g_height, 0.0],
+                },
+                Vertex {
+                    position: [left_bound+w_step/2.0, upper_bound+h_step/2.0, 0.0],
+                    world_pos: [x as f32 / g_width, y as f32 / g_height, 0.0],
+                },
+                Vertex {
+                    position: [left_bound, upper_bound+h_step, 0.0],
+                    world_pos: [x as f32 / g_width, y as f32 / g_height, 0.0],
+                }
+            ]
+        }).collect();
+        
+        let vertex_buffer = device.create_buffer_init(
+            &wgpu::util::BufferInitDescriptor {
+                label: Some("Vertex Buffer"),
+                contents: bytemuck::cast_slice(&vertices),
+                usage: wgpu::BufferUsages::VERTEX,
+            }
+        );
+       
+
+        let texture_size = wgpu::Extent3d {
+            width: domain.inner_dims().0 as u32 + 2,
+            height: domain.inner_dims().1 as u32 + 2,
+            depth_or_array_layers: 1,
+        };
+        let diffuse_texture = device.create_texture(
+            &wgpu::TextureDescriptor {
+                size: texture_size,
+                mip_level_count: 1, // We'll talk about this a little later
+                sample_count: 1,
+                dimension: wgpu::TextureDimension::D2,
+                // Most images are stored using sRGB, so we need to reflect that here.
+                format: wgpu::TextureFormat::Rgba8Unorm,
+                // TEXTURE_BINDING tells wgpu that we want to use this texture in shaders
+                // COPY_DST means that we want to copy data to this texture
+                usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
+                label: Some("diffuse_texture"),
+                // This is the same as with the SurfaceConfig. It
+                // specifies what texture formats can be used to
+                // create TextureViews for this texture. The base
+                // texture format (Rgba8UnormSrgb in this case) is
+                // always supported. Note that using a different
+                // texture format is not supported on the WebGL2
+                // backend.
+                view_formats: &[],
+            }
+        );
+
+        // We don't need to configure the texture view much, so let's
+        // let wgpu define it.
+        let diffuse_texture_view = diffuse_texture.create_view(&wgpu::TextureViewDescriptor::default());
+        let diffuse_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
+            address_mode_u: wgpu::AddressMode::ClampToEdge,
+            address_mode_v: wgpu::AddressMode::ClampToEdge,
+            address_mode_w: wgpu::AddressMode::ClampToEdge,
+            mag_filter: wgpu::FilterMode::Linear,
+            min_filter: wgpu::FilterMode::Nearest,
+            mipmap_filter: wgpu::MipmapFilterMode::Nearest,
+            ..Default::default()
+        });
+
+        let texture_bind_group_layout =
+            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+                entries: &[
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 0,
+                        visibility: wgpu::ShaderStages::FRAGMENT,
+                        ty: wgpu::BindingType::Texture {
+                            multisampled: false,
+                            view_dimension: wgpu::TextureViewDimension::D2,
+                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
+                        },
+                        count: None,
+                    },
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 1,
+                        visibility: wgpu::ShaderStages::FRAGMENT,
+                        // This should match the filterable field of the
+                        // corresponding Texture entry above.
+                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
+                        count: None,
+                    },
+                ],
+                label: Some("texture_bind_group_layout"),
+            });
+
+        let diffuse_bind_group = device.create_bind_group(
+            &wgpu::BindGroupDescriptor {
+                layout: &texture_bind_group_layout,
+                entries: &[
+                    wgpu::BindGroupEntry {
+                        binding: 0,
+                        resource: wgpu::BindingResource::TextureView(&diffuse_texture_view),
+                    },
+                    wgpu::BindGroupEntry {
+                        binding: 1,
+                        resource: wgpu::BindingResource::Sampler(&diffuse_sampler),
+                    }
+                ],
+                label: Some("diffuse_bind_group"),
+            }
+        );
+
+        let render_pipeline_layout =
+            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+                label: Some("Render Pipeline Layout"),
+                bind_group_layouts: &[&texture_bind_group_layout],
+                immediate_size: 0,
+            });
+
+        let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+            label: Some("Render Pipeline"),
+            layout: Some(&render_pipeline_layout),
+            vertex: wgpu::VertexState {
+                module: &shader,
+                entry_point: Some("vs_main"), // 1.
+                buffers: &[
+                    Vertex::desc()
+                ], // 2.
+                compilation_options: wgpu::PipelineCompilationOptions::default(),
+            },
+            fragment: Some(wgpu::FragmentState { // 3.
+                module: &shader,
+                entry_point: Some("fs_main"),
+                targets: &[Some(wgpu::ColorTargetState { // 4.
+                    format: surface_format,
+                    blend: Some(wgpu::BlendState::REPLACE),
+                    write_mask: wgpu::ColorWrites::ALL,
+                })],
+                compilation_options: wgpu::PipelineCompilationOptions::default(),
+            }),
+            // continued ...
+            primitive: wgpu::PrimitiveState {
+                topology: wgpu::PrimitiveTopology::TriangleList, // 1.
+                strip_index_format: None,
+                front_face: wgpu::FrontFace::Ccw, // 2.
+                cull_mode: Some(wgpu::Face::Back),
+                // Setting this to anything other than Fill requires Features::NON_FILL_POLYGON_MODE
+                polygon_mode: wgpu::PolygonMode::Fill,
+                // Requires Features::DEPTH_CLIP_CONTROL
+                unclipped_depth: false,
+                // Requires Features::CONSERVATIVE_RASTERIZATION
+                conservative: false,
+            },
+            // continued ...
+            depth_stencil: None, // 1.
+            multisample: wgpu::MultisampleState {
+                count: 1, // 2.
+                mask: !0, // 3.
+                alpha_to_coverage_enabled: false, // 4.
+            },
+            // multiview: None, // 5.
+            multiview_mask: None,
+            cache: None, // 6.
         });
 
 
-        for (x, y, d) in domain.smoke_iter() {
-            let inter = ((d - min_v) / (max_v - min_v)).sqrt();
-            let color = ratatui::style::Color::Rgb((255.0*inter) as u8, (255.0*inter) as u8, 0);
 
-                ctx.draw(&Points {
-                    coords: &[(-0.5 + x as f64, -0.5 + y as f64)],
-                    color: color,
-                })
-        }
-    });
-    frame.render_widget(canvas, frame.area());
+        let state = State {
+            window,
+            device,
+            queue: Arc::new(queue),
+            size,
+            surface,
+            surface_format,
+            render_pipeline,
+            vertex_buffer: Arc::new(vertex_buffer),
+            vertices,
+            texture: Arc::new(diffuse_texture),
+            diffuse_bind_group,
+        };
 
-    // let text = Text::raw("Hello World!");
-    // frame.render_widget(text, frame.area());
+        // Configure surface for the first time
+        state.configure_surface();
+
+        state
+    }
+
+    fn get_window(&self) -> &Window {
+        &self.window
+    }
+
+    fn configure_surface(&self) {
+        let surface_config = wgpu::SurfaceConfiguration {
+            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
+            format: self.surface_format,
+            // Request compatibility with the sRGB-format texture view we‘re going to create later.
+            view_formats: vec![self.surface_format.add_srgb_suffix()],
+            alpha_mode: wgpu::CompositeAlphaMode::Auto,
+            width: self.size.width,
+            height: self.size.height,
+            desired_maximum_frame_latency: 2,
+            present_mode: wgpu::PresentMode::AutoVsync,
+        };
+        self.surface.configure(&self.device, &surface_config);
+    }
+
+    fn resize(&mut self, new_size: winit::dpi::PhysicalSize<u32>) {
+        self.size = new_size;
+
+        // reconfigure the surface
+        self.configure_surface();
+    }
+
+    fn render(&mut self) {
+        // Create texture view
+        let surface_texture = self
+            .surface
+            .get_current_texture()
+            .expect("failed to acquire next swapchain texture");
+        let texture_view = surface_texture
+            .texture
+            .create_view(&wgpu::TextureViewDescriptor {
+                // Without add_srgb_suffix() the image we will be working with
+                // might not be "gamma correct".
+                format: Some(self.surface_format.add_srgb_suffix()),
+                ..Default::default()
+            });
+
+
+        // Renders a GREEN screen
+        let mut encoder = self.device.create_command_encoder(&Default::default());
+        // Create the renderpass which will clear the screen.
+        let mut renderpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+            label: None,
+            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
+                view: &texture_view,
+                depth_slice: None,
+                resolve_target: None,
+                ops: wgpu::Operations {
+                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
+                    store: wgpu::StoreOp::Store,
+                },
+            })],
+            depth_stencil_attachment: None,
+            timestamp_writes: None,
+            occlusion_query_set: None,
+            multiview_mask: None,
+        });
+
+        // If you wanted to call any drawing commands, they would go here.
+        renderpass.set_pipeline(&self.render_pipeline);
+
+        renderpass.set_bind_group(0, &self.diffuse_bind_group, &[]);
+        renderpass.set_vertex_buffer(0, self.vertex_buffer.slice(..)); 
+        renderpass.draw(0..self.vertices.len() as u32, 0..1);
+
+        // End the renderpass.
+        drop(renderpass);
+
+        // Submit the command in the queue to execute
+        self.queue.submit([encoder.finish()]);
+        self.window.pre_present_notify();
+        surface_texture.present();
+    }
+}
+
+#[derive(Default)]
+struct App {
+    state: Option<State>,
+    vis: Arc<std::sync::atomic::AtomicU8>,
+}
+
+impl ApplicationHandler for App {
+    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
+        // Create window object
+        let window = Arc::new(
+            event_loop
+                .create_window(Window::default_attributes())
+                .unwrap(),
+        );
+
+        let mut domain = fluidsim::ten_minute_physics::Domain::new(100.0, 100, 100);
+        domain.set_s_with(|x, y| {
+            if x > 30 && x < 40 && y > 40 && y < 60 {
+                return 0.0;
+            }
+
+            1.0
+        });
+        domain.set_u_with(|x, y| {
+            if x == 0 {
+                2.0
+            } else {
+                0.0
+            }
+        });
+        domain.set_smoke_with(|x, y| {
+            if x == 0 {
+                2.0
+            } else {
+                0.0
+            }
+        });
+
+
+        let state = pollster::block_on(State::new(
+            event_loop.owned_display_handle(),
+            window.clone(),
+            &domain
+        ));
+        
+        let vis = Arc::new(std::sync::atomic::AtomicU8::new(Visualisations::AbsoluteVelocity as u8));
+
+        let queue = state.queue.clone();
+        let texture = state.texture.clone();
+        let vis_clone = vis.clone();
+
+        std::thread::spawn(move || {
+            simulate(domain, queue, texture, vis_clone);
+        });
+
+        self.state = Some(state);
+        self.vis = vis;
+
+        window.request_redraw();
+    }
+
+    fn window_event(&mut self, event_loop: &ActiveEventLoop, _id: WindowId, event: WindowEvent) {
+        let state = self.state.as_mut().unwrap();
+        match event {
+            WindowEvent::CloseRequested => {
+                println!("The close button was pressed; stopping");
+                event_loop.exit();
+            }
+            WindowEvent::RedrawRequested => {
+                state.render();
+                // Emits a new redraw requested event.
+                state.get_window().request_redraw();
+            }
+            WindowEvent::Resized(size) => {
+                // Reconfigures the size of the surface. We do not re-render
+                // here as this event is always followed up by redraw request.
+                state.resize(size);
+            }
+            WindowEvent::KeyboardInput { device_id, event, is_synthetic } => {
+                match event.logical_key.as_ref() {
+                    Key::Character("1") => {
+                        self.vis.store(Visualisations::AbsoluteVelocity.into(), std::sync::atomic::Ordering::SeqCst);
+                    }
+                    Key::Character("2") => {
+                        self.vis.store(Visualisations::Smoke.into(), std::sync::atomic::Ordering::SeqCst);
+                    }
+                    Key::Character("3") => {
+                        self.vis.store(Visualisations::Pressure.into(), std::sync::atomic::Ordering::SeqCst);
+                    }
+                    _ => {}
+                };
+            }
+            _ => (),
+        }
+    }
+}
+
+#[repr(u8)]
+enum Visualisations {
+    AbsoluteVelocity,
+    Smoke,
+    Pressure,
+}
+
+impl From<u8> for Visualisations {
+    fn from(value: u8) -> Self {
+        match value {
+            0 => Self::AbsoluteVelocity,
+            1 => Self::Smoke,
+            2 => Self::Pressure,
+            other => panic!("{}", other),
+        }
+    }
+}
+impl Into<u8> for Visualisations {
+    fn into(self) -> u8 {
+        match self {
+            Self::AbsoluteVelocity => 0,
+            Self::Smoke => 1,
+            Self::Pressure => 2,
+        }
+    }
+}
+
+fn simulate(mut domain: fluidsim::ten_minute_physics::Domain, queue: Arc<wgpu::Queue>, texture: Arc<wgpu::Texture>, vis: Arc<std::sync::atomic::AtomicU8>) {
+    for i in 0.. {
+        println!("Iteration: {}", i);
+
+        domain.reset_pressure();
+
+        domain.solve_incompressibility(75, 0.01);
+        // domain.extrapolate();
+        domain.advect_vel(0.01);
+        domain.advect_smoke(0.01);
+
+        let data = match Visualisations::from(vis.load(std::sync::atomic::Ordering::SeqCst)) {
+            Visualisations::AbsoluteVelocity => {
+                let (min_div, max_div) = domain.vel_iter().fold((f32::MAX, f32::MIN), |acc, (_, _, v)| {
+                    let v = (v.0*v.0 + v.1*v.1).sqrt();
+                    (acc.0.min(v), acc.1.max(v))
+                });
+                domain.vel_iter().flat_map(|(_, _, v)| {
+                    let v = (v.0*v.0 + v.1*v.1).sqrt();
+                    let v = (v - min_div) / (max_div - min_div);
+
+                    [(v * 255.0) as u8, 0, 0, 255]
+                }).collect::<Vec<_>>()
+            }
+            Visualisations::Smoke => {
+                let (min_div, max_div) = domain.smoke_iter().fold((f32::MAX, f32::MIN), |acc, (_, _, v)| {
+                    (acc.0.min(v), acc.1.max(v))
+                });
+                domain.smoke_iter().flat_map(|(_, _, v)| {
+                    let v = (v - min_div) / (max_div - min_div);
+
+                    [(v * 255.0) as u8, 0, 0, 255]
+                }).collect::<Vec<_>>()
+            }
+            Visualisations::Pressure => {
+                let (min_div, max_div) = domain.pressure_iter().fold((f32::MAX, f32::MIN), |acc, (_, _, v)| {
+                    (acc.0.min(v), acc.1.max(v))
+                });
+                domain.pressure_iter().flat_map(|(_, _, v)| {
+                    let v = (v - min_div) / (max_div - min_div);
+
+                    [(v * 255.0) as u8, 0, 0, 255]
+                }).collect::<Vec<_>>()
+            }
+        };
+
+        let dims = domain.inner_dims();
+        let texture_size = wgpu::Extent3d {
+            width: domain.inner_dims().0 as u32 + 2,
+            height: domain.inner_dims().1 as u32 + 2,
+            depth_or_array_layers: 1,
+        };
+
+        assert_eq!(data.len() as u32, texture_size.width*texture_size.height*4);
+
+        queue.write_texture(
+            // Tells wgpu where to copy the pixel data
+            wgpu::TexelCopyTextureInfo {
+                texture: &texture,
+                mip_level: 0,
+                origin: wgpu::Origin3d::ZERO,
+                aspect: wgpu::TextureAspect::All,
+            },
+            // The actual pixel data
+            &data,
+            // The layout of the texture
+            wgpu::TexelCopyBufferLayout {
+                offset: 0,
+                bytes_per_row: Some(4 * (dims.0 as u32 + 2)),
+                rows_per_image: Some(dims.1 as u32 + 2),
+            },
+            texture_size,
+        );
+        queue.submit([]);
+        
+    }
+}
+
+fn main() {
+    // wgpu uses `log` for all of our logging, so we initialize a logger with the `env_logger` crate.
+    //
+    // To change the log level, set the `RUST_LOG` environment variable. See the `env_logger`
+    // documentation for more information.
+    env_logger::init();
+
+    let event_loop = EventLoop::new().unwrap();
+
+    // When the current loop iteration finishes, immediately begin a new
+    // iteration regardless of whether or not new events are available to
+    // process. Preferred for applications that want to render as fast as
+    // possible, like games.
+    event_loop.set_control_flow(ControlFlow::Poll);
+
+    // When the current loop iteration finishes, suspend the thread until
+    // another event arrives. Helps keeping CPU utilization low if nothing
+    // is happening, which is preferred if the application might be idling in
+    // the background.
+    // event_loop.set_control_flow(ControlFlow::Wait);
+
+    let mut app = App::default();
+    event_loop.run_app(&mut app).unwrap();
 }

src/shader.wgsl

@@ -0,0 +1,33 @@
+// Vertex shader
+
+struct VertexInput {
+    @location(0) position: vec3<f32>,
+    @location(1) world_pos: vec3<f32>,
+};
+
+struct VertexOutput {
+    @builtin(position) clip_position: vec4<f32>,
+    @location(0) world_pos: vec3<f32>,
+};
+
+@vertex
+fn vs_main(
+    model: VertexInput,
+) -> VertexOutput {
+    var out: VertexOutput;
+    out.world_pos = model.world_pos;
+    out.clip_position = vec4<f32>(model.position, 1.0);
+    return out;
+}
+
+// Fragment shader
+
+@group(0) @binding(0)
+var t_diffuse: texture_2d<f32>;
+@group(0) @binding(1)
+var s_diffuse: sampler;
+
+@fragment
+fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
+    return textureSample(t_diffuse, s_diffuse, vec2<f32>(in.world_pos[0], in.world_pos[1]));
+}

src/ten_minute_physics.rs

@@ -31,7 +31,7 @@ impl Domain {
             v: vec![0.0; num_cells],
             s: vec![0.0; num_cells],
             p: vec![0.0; num_cells],
-            m: vec![1.0; num_cells],
+            m: vec![0.0; num_cells],
         }
     }
 
@@ -60,54 +60,65 @@ impl Domain {
             }
         }
     }
+    pub fn set_smoke_with<F>(&mut self, func: F) where F: Fn(usize, usize) -> f32 {
+        for i in 1..self.width+1 {
+            for j in 1..self.height+1 {
+                self.m[i*(self.height+2)+j] = func(i-1, j-1);
+            }
+        }
+    }
+
 
     pub fn inner_dims(&self) -> (usize, usize) {
         (self.width, self.height)
     }
     
+    pub fn all_grid_iter(&self) -> impl Iterator<Item = (usize, usize)> {
+        (0..self.height+2).flat_map(|y| (0..self.width+2).zip(core::iter::repeat(y)))
+    }
     pub fn div_iter(&self) -> impl Iterator<Item = (usize, usize, f32)> {
         let n = self.height + 2;
 
-        (1..self.width+1)
-            .flat_map(|x| core::iter::repeat(x).zip(1..self.height+1))
-            .map(move |(x, y)| {
-                let div = self.u[(x+1)*n + y] - self.u[x*n + y] + self.v[x*n + y + 1] - self.v[x*n+y];
+        self.all_grid_iter().map(move |(x, y)| {
+            if x == 0 || x == self.width+1 || y == 0 || y == self.height+1 {
+                return (x, y, 0.0);
+            }
 
-                (x - 1, y - 1, div)
-            })
+            let div = self.u[(x+1)*n + y] - self.u[x*n + y] + self.v[x*n + y + 1] - self.v[x*n+y];
+
+            (x, y, div)
+        })
     }
 
-    pub fn vel_iter(&self) -> impl Iterator<Item = (usize, usize, f32)> {
+    pub fn vel_iter(&self) -> impl Iterator<Item = (usize, usize, (f32, f32))> {
         let n = self.height + 2;
 
-        (1..self.width+1)
-            .flat_map(|x| core::iter::repeat(x).zip(1..self.height+1))
-            .map(move |(x, y)| {
-                let u = self.u[x*n+y];
-                let v = self.v[x*n+y];
+        self.all_grid_iter().map(move |(x ,y)| {
+            if self.s[x*n+y] == 0.0 {
+                return (x, y, (0.0, 0.0));
+            }
 
-                (x - 1, y - 1, (u*u + v*v).sqrt())
-            })
+            let u = self.u[x*n+y];
+            let v = self.v[x*n+y];
+
+            (x, y, (u, v))
+        })
     }
 
     pub fn pressure_iter(&self) -> impl Iterator<Item = (usize, usize, f32)> {
         let n = self.height + 2;
 
-        (1..self.width+1)
-            .flat_map(|x| core::iter::repeat(x).zip(1..self.height+1))
-            .map(move |(i, j)| {
-                (i - 1, j - 1, self.p[i*n+j])
-            })
+        self.all_grid_iter().map(move |(x, y)| {
+            (x, y, self.p[x*n + y])
+        })
     }
 
     pub fn smoke_iter(&self) -> impl Iterator<Item = (usize, usize, f32)> {
         let n = self.height + 2;
 
-        (1..self.width+1)
-            .flat_map(|x| core::iter::repeat(x).zip(1..self.height+1))
-            .map(move |(i, j)| {
-                (i - 1, j - 1, self.m[i*n+j])
-            })
+        self.all_grid_iter().map(move |(x, y)| {
+            (x, y, self.m[x*n+y])
+        })
     }
 
     pub fn solve_incompressibility(&mut self, num_iters: usize, dT: f32) {
@@ -135,7 +146,7 @@ impl Domain {
                     let div = self.u[(i+1)*n + j] - self.u[i*n + j] + self.v[i*n + j + 1] - self.v[i*n+j];
                     
                     let p = -div / s;
-                    let p = p *  1.9;
+                    // let p = p *  1.4;
                     self.p[i*n+j] += cp * p;
 
                     self.u[i*n + j] -= sx0 * p;