Skip to main content

pfx_physics/
ripple.rs

1use bytemuck::{Pod, Zeroable};
2use half::f16;
3
4pub const DAMP: f32 = 0.991;
5pub const DROPS: usize = 8;
6pub const FIXED_DT: f32 = 1.0 / 60.0;
7pub const RIPPLE_WGSL: &str = include_str!("ripple.wgsl");
8pub const CAUSTICS_WGSL: &str = include_str!("caustics.wgsl");
9pub const CAUSTIC_GRID: [u32; 2] = [640, 400];
10
11fn caustic_target(device: &wgpu::Device, frame_size: [u32; 2]) -> wgpu::Texture {
12    device.create_texture(&wgpu::TextureDescriptor {
13        label: Some("Caustics"),
14        size: wgpu::Extent3d {
15            width: (frame_size[0] / 2).max(1),
16            height: (frame_size[1] / 2).max(1),
17            depth_or_array_layers: 1,
18        },
19        mip_level_count: 1,
20        sample_count: 1,
21        dimension: wgpu::TextureDimension::D2,
22        format: wgpu::TextureFormat::Rgba16Float,
23        usage: wgpu::TextureUsages::RENDER_ATTACHMENT
24            | wgpu::TextureUsages::TEXTURE_BINDING
25            | wgpu::TextureUsages::COPY_SRC,
26        view_formats: &[],
27    })
28}
29
30#[repr(C)]
31#[derive(Clone, Copy, Pod, Zeroable)]
32struct CausticParams {
33    size: [f32; 2],
34    time: f32,
35    calm: f32,
36}
37
38#[repr(C)]
39#[derive(Clone, Copy, Pod, Zeroable, Debug)]
40pub struct Drop {
41    pub uv: [f32; 2],
42    pub radius: f32,
43    pub amount: f32,
44}
45
46#[repr(C)]
47#[derive(Clone, Copy, Pod, Zeroable)]
48struct DropBlock {
49    head: [f32; 4],
50    items: [[f32; 4]; DROPS],
51}
52
53pub struct Ripple {
54    width: u32,
55    height: u32,
56    now: Vec<f32>,
57    prev: Vec<f32>,
58    damp: f32,
59    fixed_dt: f32,
60    clock: f32,
61}
62
63impl Ripple {
64    pub fn new(width: u32, height: u32) -> Self {
65        Self::with_damp(width, height, DAMP)
66    }
67
68    pub fn undamped(width: u32, height: u32) -> Self {
69        Self::with_damp(width, height, 1.0)
70    }
71
72    fn with_damp(width: u32, height: u32, damp: f32) -> Self {
73        let width = width.max(1);
74        let height = height.max(1);
75        let n = (width * height) as usize;
76        Self {
77            width,
78            height,
79            now: vec![0.0; n],
80            prev: vec![0.0; n],
81            damp,
82            fixed_dt: FIXED_DT,
83            clock: 0.0,
84        }
85    }
86
87    pub fn width(&self) -> u32 {
88        self.width
89    }
90
91    pub fn height(&self) -> u32 {
92        self.height
93    }
94
95    pub fn heights(&self) -> &[f32] {
96        &self.now
97    }
98
99    pub fn previous(&self) -> &[f32] {
100        &self.prev
101    }
102
103    pub fn damp(&self) -> f32 {
104        self.damp
105    }
106
107    pub fn write(&mut self, x: u32, y: u32, height: f32) {
108        if let Some(slot) = self.index(x, y) {
109            self.now[slot] = height;
110            self.prev[slot] = height;
111        }
112    }
113
114    pub fn height_at(&self, x: i32, y: i32) -> f32 {
115        let x = x.clamp(0, self.width as i32 - 1) as u32;
116        let y = y.clamp(0, self.height as i32 - 1) as u32;
117        self.now[(y * self.width + x) as usize]
118    }
119
120    pub fn step(&mut self, dt: f32, drops: &[Drop]) {
121        if !dt.is_finite() || dt <= 0.0 {
122            return;
123        }
124        self.clock = (self.clock + dt).min(self.fixed_dt * 8.0);
125        let steps = (self.clock / self.fixed_dt).floor() as usize;
126        self.clock -= steps as f32 * self.fixed_dt;
127        for step in 0..steps {
128            self.iterate(if step == 0 { drops } else { &[] });
129        }
130    }
131
132    fn index(&self, x: u32, y: u32) -> Option<usize> {
133        if x < self.width && y < self.height {
134            Some((y * self.width + x) as usize)
135        } else {
136            None
137        }
138    }
139
140    fn at(&self, x: i32, y: i32) -> f32 {
141        let x = x.clamp(0, self.width as i32 - 1) as u32;
142        let y = y.clamp(0, self.height as i32 - 1) as u32;
143        self.now[(y * self.width + x) as usize]
144    }
145
146    fn iterate(&mut self, drops: &[Drop]) {
147        let w = self.width as i32;
148        let h = self.height as i32;
149        let mut next = vec![0.0; self.now.len()];
150        for y in 0..h {
151            for x in 0..w {
152                let around =
153                    self.at(x - 1, y) + self.at(x + 1, y) + self.at(x, y - 1) + self.at(x, y + 1);
154                let prev = self.prev[(y as u32 * self.width + x as u32) as usize];
155                let mut height = (around * 0.5 - prev) * self.damp;
156                let uv = [(x as f32 + 0.5) / w as f32, (y as f32 + 0.5) / h as f32];
157                let aspect = [w as f32 / h as f32, 1.0];
158                for drop in drops.iter().take(DROPS) {
159                    let q = [
160                        (uv[0] - drop.uv[0]) * aspect[0] / drop.radius.max(1e-4),
161                        (uv[1] - drop.uv[1]) * aspect[1] / drop.radius.max(1e-4),
162                    ];
163                    height += drop.amount * (-(q[0] * q[0] + q[1] * q[1])).exp();
164                }
165                next[(y as u32 * self.width + x as u32) as usize] = height;
166            }
167        }
168        self.prev.clone_from(&self.now);
169        self.now = next;
170    }
171}
172
173pub struct RippleGpu {
174    width: u32,
175    height: u32,
176    textures: [wgpu::Texture; 2],
177    current: usize,
178    drops: wgpu::Buffer,
179    pipeline: wgpu::ComputePipeline,
180    binds: [[wgpu::BindGroup; 2]; 2],
181    caustic_texture: wgpu::Texture,
182    caustic_pipeline: wgpu::RenderPipeline,
183    caustic_binds: [wgpu::BindGroup; 2],
184    caustic_params: wgpu::Buffer,
185    frame_size: [u32; 2],
186    fixed_dt: f32,
187    clock: f32,
188}
189
190impl RippleGpu {
191    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue, width: u32, height: u32) -> Self {
192        let width = width.max(1);
193        let height = height.max(1);
194        let texture = |label: &str| {
195            device.create_texture(&wgpu::TextureDescriptor {
196                label: Some(label),
197                size: wgpu::Extent3d {
198                    width,
199                    height,
200                    depth_or_array_layers: 1,
201                },
202                mip_level_count: 1,
203                sample_count: 1,
204                dimension: wgpu::TextureDimension::D2,
205                format: wgpu::TextureFormat::Rgba16Float,
206                usage: wgpu::TextureUsages::TEXTURE_BINDING
207                    | wgpu::TextureUsages::STORAGE_BINDING
208                    | wgpu::TextureUsages::COPY_SRC
209                    | wgpu::TextureUsages::RENDER_ATTACHMENT,
210                view_formats: &[],
211            })
212        };
213        let textures = [texture("Ripple A"), texture("Ripple B")];
214        let views = [
215            textures[0].create_view(&Default::default()),
216            textures[1].create_view(&Default::default()),
217        ];
218        let block = |label: &str| {
219            device.create_buffer(&wgpu::BufferDescriptor {
220                label: Some(label),
221                size: std::mem::size_of::<DropBlock>() as u64,
222                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
223                mapped_at_creation: false,
224            })
225        };
226        let drops = block("Ripple drops");
227        let none = block("Ripple no drops");
228        queue.write_buffer(&none, 0, bytemuck::bytes_of(&block_of(&[])));
229        let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
230            label: Some("Ripple"),
231            entries: &[
232                wgpu::BindGroupLayoutEntry {
233                    binding: 0,
234                    visibility: wgpu::ShaderStages::COMPUTE,
235                    ty: wgpu::BindingType::Texture {
236                        sample_type: wgpu::TextureSampleType::Float { filterable: false },
237                        view_dimension: wgpu::TextureViewDimension::D2,
238                        multisampled: false,
239                    },
240                    count: None,
241                },
242                wgpu::BindGroupLayoutEntry {
243                    binding: 1,
244                    visibility: wgpu::ShaderStages::COMPUTE,
245                    ty: wgpu::BindingType::StorageTexture {
246                        access: wgpu::StorageTextureAccess::WriteOnly,
247                        format: wgpu::TextureFormat::Rgba16Float,
248                        view_dimension: wgpu::TextureViewDimension::D2,
249                    },
250                    count: None,
251                },
252                wgpu::BindGroupLayoutEntry {
253                    binding: 2,
254                    visibility: wgpu::ShaderStages::COMPUTE,
255                    ty: wgpu::BindingType::Buffer {
256                        ty: wgpu::BufferBindingType::Uniform,
257                        has_dynamic_offset: false,
258                        min_binding_size: None,
259                    },
260                    count: None,
261                },
262            ],
263        });
264        let bind = |src: usize, dst: usize, input: &wgpu::Buffer| {
265            device.create_bind_group(&wgpu::BindGroupDescriptor {
266                label: Some("Ripple"),
267                layout: &layout,
268                entries: &[
269                    wgpu::BindGroupEntry {
270                        binding: 0,
271                        resource: wgpu::BindingResource::TextureView(&views[src]),
272                    },
273                    wgpu::BindGroupEntry {
274                        binding: 1,
275                        resource: wgpu::BindingResource::TextureView(&views[dst]),
276                    },
277                    wgpu::BindGroupEntry {
278                        binding: 2,
279                        resource: input.as_entire_binding(),
280                    },
281                ],
282            })
283        };
284        let binds = [
285            [bind(0, 1, &drops), bind(0, 1, &none)],
286            [bind(1, 0, &drops), bind(1, 0, &none)],
287        ];
288        let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
289            label: Some("Ripple"),
290            source: wgpu::ShaderSource::Wgsl(RIPPLE_WGSL.into()),
291        });
292        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
293            label: Some("Ripple"),
294            layout: Some(
295                &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
296                    label: Some("Ripple"),
297                    bind_group_layouts: &[Some(&layout)],
298                    immediate_size: 0,
299                }),
300            ),
301            module: &module,
302            entry_point: Some("ripple"),
303            compilation_options: Default::default(),
304            cache: None,
305        });
306        let caustic_params = device.create_buffer(&wgpu::BufferDescriptor {
307            label: Some("Caustic parameters"),
308            size: std::mem::size_of::<CausticParams>() as u64,
309            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
310            mapped_at_creation: false,
311        });
312        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
313            label: Some("Caustic ripple sampler"),
314            mag_filter: wgpu::FilterMode::Linear,
315            min_filter: wgpu::FilterMode::Linear,
316            ..Default::default()
317        });
318        let caustic_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
319            label: Some("Caustics"),
320            entries: &[
321                wgpu::BindGroupLayoutEntry {
322                    binding: 0,
323                    visibility: wgpu::ShaderStages::VERTEX,
324                    ty: wgpu::BindingType::Texture {
325                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
326                        view_dimension: wgpu::TextureViewDimension::D2,
327                        multisampled: false,
328                    },
329                    count: None,
330                },
331                wgpu::BindGroupLayoutEntry {
332                    binding: 1,
333                    visibility: wgpu::ShaderStages::VERTEX,
334                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
335                    count: None,
336                },
337                wgpu::BindGroupLayoutEntry {
338                    binding: 2,
339                    visibility: wgpu::ShaderStages::VERTEX,
340                    ty: wgpu::BindingType::Buffer {
341                        ty: wgpu::BufferBindingType::Uniform,
342                        has_dynamic_offset: false,
343                        min_binding_size: None,
344                    },
345                    count: None,
346                },
347            ],
348        });
349        let caustic_binds = std::array::from_fn(|i| {
350            device.create_bind_group(&wgpu::BindGroupDescriptor {
351                label: Some("Caustics"),
352                layout: &caustic_layout,
353                entries: &[
354                    wgpu::BindGroupEntry {
355                        binding: 0,
356                        resource: wgpu::BindingResource::TextureView(&views[i]),
357                    },
358                    wgpu::BindGroupEntry {
359                        binding: 1,
360                        resource: wgpu::BindingResource::Sampler(&sampler),
361                    },
362                    wgpu::BindGroupEntry {
363                        binding: 2,
364                        resource: caustic_params.as_entire_binding(),
365                    },
366                ],
367            })
368        });
369        let caustic_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
370            label: Some("Caustics"),
371            source: wgpu::ShaderSource::Wgsl(CAUSTICS_WGSL.into()),
372        });
373        let caustic_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
374            label: Some("Caustics"),
375            layout: Some(
376                &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
377                    label: Some("Caustics"),
378                    bind_group_layouts: &[Some(&caustic_layout)],
379                    immediate_size: 0,
380                }),
381            ),
382            vertex: wgpu::VertexState {
383                module: &caustic_module,
384                entry_point: Some("vs_caustic"),
385                compilation_options: Default::default(),
386                buffers: &[],
387            },
388            primitive: wgpu::PrimitiveState::default(),
389            depth_stencil: None,
390            multisample: wgpu::MultisampleState::default(),
391            fragment: Some(wgpu::FragmentState {
392                module: &caustic_module,
393                entry_point: Some("fs_caustic"),
394                compilation_options: Default::default(),
395                targets: &[Some(wgpu::ColorTargetState {
396                    format: wgpu::TextureFormat::Rgba16Float,
397                    blend: Some(wgpu::BlendState {
398                        color: wgpu::BlendComponent {
399                            src_factor: wgpu::BlendFactor::One,
400                            dst_factor: wgpu::BlendFactor::One,
401                            operation: wgpu::BlendOperation::Add,
402                        },
403                        alpha: wgpu::BlendComponent::REPLACE,
404                    }),
405                    write_mask: wgpu::ColorWrites::ALL,
406                })],
407            }),
408            multiview_mask: None,
409            cache: None,
410        });
411        let frame_size = [width, height];
412        let caustic_texture = caustic_target(device, frame_size);
413        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
414            label: Some("Ripple clear"),
415        });
416        for view in &views {
417            let pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
418                multiview_mask: None,
419                label: Some("Ripple clear"),
420                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
421                    view,
422                    depth_slice: None,
423                    resolve_target: None,
424                    ops: wgpu::Operations {
425                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
426                        store: wgpu::StoreOp::Store,
427                    },
428                })],
429                depth_stencil_attachment: None,
430                timestamp_writes: None,
431                occlusion_query_set: None,
432            });
433            drop(pass);
434        }
435        queue.submit(Some(encoder.finish()));
436        Self {
437            width,
438            height,
439            textures,
440            current: 0,
441            drops,
442            pipeline,
443            binds,
444            caustic_texture,
445            caustic_pipeline,
446            caustic_binds,
447            caustic_params,
448            frame_size,
449            fixed_dt: FIXED_DT,
450            clock: 0.0,
451        }
452    }
453
454    pub fn texture(&self) -> &wgpu::Texture {
455        &self.textures[self.current]
456    }
457
458    pub fn caustics(&self) -> &wgpu::Texture {
459        &self.caustic_texture
460    }
461
462    pub fn resize_caustics(&mut self, device: &wgpu::Device, width: u32, height: u32) {
463        let frame_size = [width.max(1), height.max(1)];
464        if frame_size != self.frame_size {
465            self.caustic_texture = caustic_target(device, frame_size);
466            self.frame_size = frame_size;
467        }
468    }
469
470    pub fn render_caustics(
471        &self,
472        device: &wgpu::Device,
473        queue: &wgpu::Queue,
474        time: f32,
475        calm: f32,
476    ) {
477        queue.write_buffer(
478            &self.caustic_params,
479            0,
480            bytemuck::bytes_of(&CausticParams {
481                size: [self.frame_size[0] as f32, self.frame_size[1] as f32],
482                time,
483                calm,
484            }),
485        );
486        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
487            label: Some("Caustics"),
488        });
489        let view = self.caustic_texture.create_view(&Default::default());
490        {
491            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
492                multiview_mask: None,
493                label: Some("Caustics"),
494                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
495                    view: &view,
496                    depth_slice: None,
497                    resolve_target: None,
498                    ops: wgpu::Operations {
499                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
500                        store: wgpu::StoreOp::Store,
501                    },
502                })],
503                depth_stencil_attachment: None,
504                timestamp_writes: None,
505                occlusion_query_set: None,
506            });
507            pass.set_pipeline(&self.caustic_pipeline);
508            pass.set_bind_group(0, &self.caustic_binds[self.current], &[]);
509            pass.draw(0..CAUSTIC_GRID[0] * CAUSTIC_GRID[1] * 6, 0..3);
510        }
511        queue.submit(Some(encoder.finish()));
512    }
513
514    pub fn read_caustics(&self, device: &wgpu::Device, queue: &wgpu::Queue) -> Vec<[f16; 4]> {
515        let size = self.caustic_texture.size();
516        let padded = (size.width * 8).div_ceil(256) * 256;
517        let staging = device.create_buffer(&wgpu::BufferDescriptor {
518            label: Some("Caustic readback"),
519            size: u64::from(padded) * u64::from(size.height),
520            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
521            mapped_at_creation: false,
522        });
523        let mut encoder = device.create_command_encoder(&Default::default());
524        encoder.copy_texture_to_buffer(
525            self.caustic_texture.as_image_copy(),
526            wgpu::TexelCopyBufferInfo {
527                buffer: &staging,
528                layout: wgpu::TexelCopyBufferLayout {
529                    offset: 0,
530                    bytes_per_row: Some(padded),
531                    rows_per_image: Some(size.height),
532                },
533            },
534            size,
535        );
536        queue.submit(Some(encoder.finish()));
537        let slice = staging.slice(..);
538        let (send, recv) = std::sync::mpsc::channel();
539        slice.map_async(wgpu::MapMode::Read, move |result| {
540            let _ = send.send(result);
541        });
542        device
543            .poll(wgpu::PollType::wait_indefinitely())
544            .expect("caustic poll");
545        recv.recv()
546            .expect("caustic map")
547            .expect("caustic buffer map");
548        let data = slice.get_mapped_range().expect("a mapped range");
549        let mut pixels = Vec::with_capacity((size.width * size.height) as usize);
550        for y in 0..size.height as usize {
551            for x in 0..size.width as usize {
552                let at = y * padded as usize + x * 8;
553                pixels.push(std::array::from_fn(|channel| {
554                    let i = at + channel * 2;
555                    f16::from_bits(u16::from_le_bytes([data[i], data[i + 1]]))
556                }));
557            }
558        }
559        drop(data);
560        staging.unmap();
561        pixels
562    }
563
564    pub fn step(&mut self, device: &wgpu::Device, queue: &wgpu::Queue, dt: f32, drops: &[Drop]) {
565        if !dt.is_finite() || dt <= 0.0 {
566            return;
567        }
568        self.clock = (self.clock + dt).min(self.fixed_dt * 8.0);
569        let steps = (self.clock / self.fixed_dt).floor() as usize;
570        self.clock -= steps as f32 * self.fixed_dt;
571        if steps == 0 {
572            return;
573        }
574        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
575            label: Some("Ripple"),
576        });
577        queue.write_buffer(&self.drops, 0, bytemuck::bytes_of(&block_of(drops)));
578        for step in 0..steps {
579            {
580                let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
581                    label: Some("Ripple"),
582                    timestamp_writes: None,
583                });
584                pass.set_pipeline(&self.pipeline);
585                pass.set_bind_group(0, &self.binds[self.current][usize::from(step > 0)], &[]);
586                pass.dispatch_workgroups(self.width.div_ceil(8), self.height.div_ceil(8), 1);
587            }
588            self.current = 1 - self.current;
589        }
590        queue.submit(Some(encoder.finish()));
591    }
592}
593
594fn block_of(drops: &[Drop]) -> DropBlock {
595    let mut block = DropBlock {
596        head: [drops.len().min(DROPS) as f32, 0.0, 0.0, 0.0],
597        items: [[0.0; 4]; DROPS],
598    };
599    for (slot, drop) in block.items.iter_mut().zip(drops.iter().take(DROPS)) {
600        *slot = [drop.uv[0], drop.uv[1], drop.radius, drop.amount];
601    }
602    block
603}
604
605#[cfg(test)]
606mod tests {
607    use super::*;
608
609    fn sample_height(heights: &[f32], size: [usize; 2], uv: [f32; 2]) -> f32 {
610        let x = uv[0] * size[0] as f32 - 0.5;
611        let y = uv[1] * size[1] as f32 - 0.5;
612        let ix = x.floor() as i32;
613        let iy = y.floor() as i32;
614        let fx = x - ix as f32;
615        let fy = y - iy as f32;
616        let at = |dx: i32, dy: i32| {
617            let x = (ix + dx).clamp(0, size[0] as i32 - 1) as usize;
618            let y = (iy + dy).clamp(0, size[1] as i32 - 1) as usize;
619            heights[y * size[0] + x]
620        };
621        let a = at(0, 0) * (1.0 - fx) + at(1, 0) * fx;
622        let b = at(0, 1) * (1.0 - fx) + at(1, 1) * fx;
623        a * (1.0 - fy) + b * fy
624    }
625
626    fn caustic_vertex(
627        heights: &[f32],
628        water_size: [usize; 2],
629        frame_size: [f32; 2],
630        cell: [u32; 2],
631        channel: usize,
632        time: f32,
633        calm: f32,
634    ) -> ([f32; 2], [f32; 2]) {
635        let uv = [
636            cell[0] as f32 / CAUSTIC_GRID[0] as f32 * 1.12 - 0.06,
637            cell[1] as f32 / CAUSTIC_GRID[1] as f32 * 1.12 - 0.06,
638        ];
639        let texel = [1.0 / water_size[0] as f32, 1.0 / water_size[1] as f32];
640        let l = sample_height(heights, water_size, [uv[0] - texel[0], uv[1]]);
641        let r = sample_height(heights, water_size, [uv[0] + texel[0], uv[1]]);
642        let d = sample_height(heights, water_size, [uv[0], uv[1] - texel[1]]);
643        let u = sample_height(heights, water_size, [uv[0], uv[1] + texel[1]]);
644        let aspect = frame_size[0] / frame_size[1].max(1.0);
645        let c = 0.62f32.cos();
646        let s = 0.62f32.sin();
647        let at = [uv[0] * aspect, uv[1]];
648        let p = [
649            (c * at[0] - s * at[1]) * 0.55,
650            (s * at[0] + c * at[1]) * 1.35,
651        ];
652        let waves = [
653            [0.83, 0.56, 14.0, 1.122],
654            [-0.44, 0.90, 19.0, 1.308],
655            [0.97, -0.24, 26.0, 1.530],
656            [-0.71, -0.70, 33.0, 1.723],
657            [0.26, 0.97, 41.0, 1.921],
658            [-0.93, 0.37, 53.0, 2.184],
659            [0.60, -0.80, 67.0, 2.456],
660            [-0.15, -0.99, 83.0, 2.733],
661        ];
662        let amps = [
663            0.00829, 0.00450, 0.00241, 0.00150, 0.00096, 0.00057, 0.00036, 0.00023,
664        ];
665        let mut grad = [0.0; 2];
666        for (wave, amp) in waves.into_iter().zip(amps) {
667            let k = wave[2] * 1.25;
668            let amplitude = amp / 1.25;
669            let phase = (wave[0] * p[0] + wave[1] * p[1]) * k + time * calm * wave[3];
670            grad[0] += amplitude * phase.cos() * wave[0] * k;
671            grad[1] += amplitude * phase.cos() * wave[1] * k;
672        }
673        let g = [grad[0] * 0.55, grad[1] * 1.35];
674        let swell = [c * g[0] + s * g[1], -s * g[0] + c * g[1]];
675        let slope = [
676            -((r - l) * 0.9 + swell[0]),
677            -((u - d) * 0.9 + swell[1]),
678            1.0,
679        ];
680        let length = (slope[0] * slope[0] + slope[1] * slope[1] + 1.0).sqrt();
681        let n = [slope[0] / length, slope[1] / length, 1.0 / length];
682        let eta = [0.736, 0.752, 0.770][channel];
683        let dot = -n[2];
684        let k = 1.0 - eta * eta * (1.0 - dot * dot);
685        let factor = eta * dot + k.sqrt();
686        let ray = [-factor * n[0], -factor * n[1], -eta - factor * n[2]];
687        let landed = [
688            uv[0] + ray[0] / (-ray[2]).max(0.2) * 0.66 / aspect,
689            uv[1] + ray[1] / (-ray[2]).max(0.2) * 0.66,
690        ];
691        (uv, landed)
692    }
693
694    fn caustic_pixel(
695        heights: &[f32],
696        water_size: [usize; 2],
697        output: [u32; 2],
698        pixel: [u32; 2],
699        time: f32,
700        calm: f32,
701    ) -> [f32; 4] {
702        let frame = [output[0] as f32 * 2.0, output[1] as f32 * 2.0];
703        let mut color = [0.0, 0.0, 0.0, 1.0];
704        let at = [pixel[0] as f32 + 0.5, pixel[1] as f32 + 0.5];
705        let near = [
706            ((at[0] / output[0] as f32 + 0.06) / 1.12 * CAUSTIC_GRID[0] as f32) as i32,
707            ((at[1] / output[1] as f32 + 0.06) / 1.12 * CAUSTIC_GRID[1] as f32) as i32,
708        ];
709        for (channel, slot) in color.iter_mut().enumerate().take(3) {
710            for cy in (near[1] - 12).max(0)..=(near[1] + 12).min(CAUSTIC_GRID[1] as i32 - 1) {
711                for cx in (near[0] - 12).max(0)..=(near[0] + 12).min(CAUSTIC_GRID[0] as i32 - 1) {
712                    let corners = [[0, 0], [1, 0], [0, 1], [1, 1]].map(|corner| {
713                        caustic_vertex(
714                            heights,
715                            water_size,
716                            frame,
717                            [cx as u32 + corner[0], cy as u32 + corner[1]],
718                            channel,
719                            time,
720                            calm,
721                        )
722                    });
723                    for indices in [[0, 1, 2], [2, 1, 3]] {
724                        let v = indices.map(|i| corners[i]);
725                        let p = v.map(|(_, after)| {
726                            [after[0] * output[0] as f32, after[1] * output[1] as f32]
727                        });
728                        let cross = |a: [f32; 2], b: [f32; 2]| a[0] * b[1] - a[1] * b[0];
729                        let e1 = [p[1][0] - p[0][0], p[1][1] - p[0][1]];
730                        let e2 = [p[2][0] - p[0][0], p[2][1] - p[0][1]];
731                        let area = cross(e1, e2);
732                        if area.abs() < 1e-9 {
733                            continue;
734                        }
735                        let q = [at[0] - p[0][0], at[1] - p[0][1]];
736                        let b = cross(q, e2) / area;
737                        let c = cross(e1, q) / area;
738                        if b < 0.0 || c < 0.0 || b + c > 1.0 {
739                            continue;
740                        }
741                        let du = [v[1].0[0] - v[0].0[0], v[1].0[1] - v[0].0[1]];
742                        let dv = [v[2].0[0] - v[0].0[0], v[2].0[1] - v[0].0[1]];
743                        let dx = [
744                            (du[0] * e2[1] - dv[0] * e1[1]) / area,
745                            (du[1] * e2[1] - dv[1] * e1[1]) / area,
746                        ];
747                        let dy = [
748                            (dv[0] * e1[0] - du[0] * e2[0]) / area,
749                            (dv[1] * e1[0] - du[1] * e2[0]) / area,
750                        ];
751                        let before = (dx[0] * dx[0] + dx[1] * dx[1]).sqrt()
752                            * (dy[0] * dy[0] + dy[1] * dy[1]).sqrt();
753                        let light =
754                            (before * (output[0] * output[1]) as f32).clamp(0.0, 14.0) * 0.34;
755                        *slot += light;
756                    }
757                }
758            }
759        }
760        color
761    }
762
763    #[test]
764    fn caustic_cpu_port_maps_a_small_ripple_grid_into_three_channels() {
765        let mut ripple = Ripple::new(8, 8);
766        ripple.step(
767            FIXED_DT,
768            &[Drop {
769                uv: [0.5, 0.5],
770                radius: 0.2,
771                amount: 0.04,
772            }],
773        );
774        let heights: Vec<f32> = ripple
775            .heights()
776            .iter()
777            .map(|&h| f16::from_f32(h).to_f32())
778            .collect();
779        let uv = caustic_vertex(&heights, [8, 8], [1280.0, 800.0], [320, 200], 0, 0.25, 1.0);
780        assert!((uv.0[0] - 0.5).abs() < 1e-6);
781        assert!((uv.0[1] - 0.5).abs() < 1e-6);
782        let red = caustic_pixel(&heights, [8, 8], [640, 400], [320, 200], 0.25, 1.0);
783        assert!(red[..3].iter().all(|v| v.is_finite() && *v > 0.0));
784        assert!((red[0] - red[2]).abs() > 0.0001);
785        assert_eq!(red[3], 1.0);
786    }
787
788    #[test]
789    fn caustic_shader_parses() {
790        naga::front::wgsl::parse_str(CAUSTICS_WGSL).unwrap();
791    }
792
793    #[test]
794    #[ignore = "needs a GPU; run under pgpu"]
795    fn caustics_after_one_drop_match_the_cpu_port() {
796        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
797            backends: wgpu::Backends::VULKAN,
798            ..wgpu::InstanceDescriptor::new_without_display_handle()
799        });
800        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
801            power_preference: wgpu::PowerPreference::HighPerformance,
802            compatible_surface: None,
803            force_fallback_adapter: false,
804            apply_limit_buckets: false,
805        }))
806        .expect("GPU adapter");
807        let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
808            label: Some("caustic test"),
809            required_features: wgpu::Features::empty(),
810            ..Default::default()
811        }))
812        .expect("GPU device");
813        let drop = Drop {
814            uv: [0.5, 0.5],
815            radius: 0.2,
816            amount: 0.04,
817        };
818        let mut cpu = Ripple::new(8, 8);
819        cpu.step(FIXED_DT, &[drop]);
820        let heights: Vec<f32> = cpu
821            .heights()
822            .iter()
823            .map(|&h| f16::from_f32(h).to_f32())
824            .collect();
825        let mut gpu = RippleGpu::new(&device, &queue, 8, 8);
826        gpu.resize_caustics(&device, 1280, 800);
827        gpu.step(&device, &queue, FIXED_DT, &[drop]);
828        gpu.render_caustics(&device, &queue, 0.25, 1.0);
829        let actual = gpu.read_caustics(&device, &queue);
830        assert_eq!(actual.len(), 640 * 400);
831        for pixel in [[320, 200], [200, 160], [430, 260]] {
832            let expected = caustic_pixel(&heights, [8, 8], [640, 400], pixel, 0.25, 1.0);
833            let seen = actual[pixel[1] as usize * 640 + pixel[0] as usize];
834            for channel in 0..3 {
835                let difference = (seen[channel].to_f32() - expected[channel]).abs();
836                assert!(
837                    difference <= 2.0 * f16::EPSILON.to_f32() * expected[channel].max(1.0),
838                    "pixel {pixel:?} channel {channel}: gpu {} cpu {}",
839                    seen[channel].to_f32(),
840                    expected[channel]
841                );
842            }
843        }
844    }
845
846    fn energy(now: &[f32], prev: &[f32], width: usize) -> f32 {
847        let height = now.len() / width;
848        let mut total = 0.0;
849        for y in 0..height {
850            for x in 0..width {
851                let i = y * width + x;
852                let d = now[i] - prev[i];
853                total += d * d;
854                if x + 1 < width {
855                    let g = now[i] - now[i + 1];
856                    total += 0.25 * g * g;
857                }
858                if y + 1 < height {
859                    let g = now[i] - now[i + width];
860                    total += 0.25 * g * g;
861                }
862            }
863        }
864        total
865    }
866
867    #[test]
868    fn one_step_matches_the_wave_stencil() {
869        let mut ripple = Ripple::new(3, 3);
870        ripple.write(1, 1, 1.0);
871        ripple.step(FIXED_DT, &[]);
872        assert_eq!(ripple.height_at(1, 1), -DAMP);
873        assert_eq!(ripple.height_at(0, 1), 0.5 * DAMP);
874        assert_eq!(ripple.height_at(0, 0), 0.0);
875        assert_eq!(ripple.previous()[3 + 1], 1.0);
876    }
877
878    #[test]
879    fn a_drop_raises_the_centre() {
880        let mut ripple = Ripple::new(16, 16);
881        ripple.step(
882            FIXED_DT,
883            &[Drop {
884                uv: [0.5, 0.5],
885                radius: 0.08,
886                amount: 1.0,
887            }],
888        );
889        assert!(ripple.height_at(8, 8) > ripple.height_at(0, 0));
890        assert!(ripple.height_at(8, 8) > 0.5);
891    }
892
893    #[test]
894    fn energy_holds_roughly_and_damping_spends_it() {
895        let mut calm = Ripple::undamped(48, 48);
896        let mut damped = Ripple::new(48, 48);
897        for y in 0..48 {
898            for x in 0..48 {
899                let dx = x as f32 - 24.0;
900                let dy = y as f32 - 24.0;
901                let h = (-(dx * dx + dy * dy) / 36.0).exp();
902                calm.write(x, y, h);
903                damped.write(x, y, h);
904            }
905        }
906        let start = energy(calm.heights(), calm.previous(), 48);
907        for _ in 0..180 {
908            calm.step(FIXED_DT, &[]);
909            damped.step(FIXED_DT, &[]);
910        }
911        let held = energy(calm.heights(), calm.previous(), 48);
912        let spent = energy(damped.heights(), damped.previous(), 48);
913        let ratio = held / start;
914        assert!(ratio > 0.5 && ratio < 2.0, "undamped ratio {ratio}");
915        assert!(spent < held * 0.5, "damped {spent} undamped {held}");
916        assert!(spent < start);
917    }
918}