Skip to main content

pfx_bake/
trace.rs

1use pfx_gpu::pace::{Pacer, Slice, Stats, Turns, Work};
2use pfx_gpu::{Gpu, wgpu};
3use pfx_load::Sky;
4use pfx_materials::Material;
5use pfx_trace::bvh::{Bvh, Hit, Ray, Triangle};
6use pfx_trace::detail::Detail;
7use pfx_trace::shapes::{Shape, intersect};
8use pfx_trace::{Camera, Sampler, Scene, Sun, Trace};
9use serde::{Deserialize, Serialize};
10
11use crate::batch::{Batch, FACES_PER_BATCH, FaceRecord};
12use crate::grid::{AXES, Grid, GridSpec, Probe, ProbeExtra};
13
14const FACE: usize = 4;
15const COSINE_NORMAL: f32 = std::f32::consts::PI / 3.253_611_3;
16const INSIDE: u16 = 6;
17const WELD: f64 = 1e5;
18const ROUNDS: u32 = 4;
19const MAX_ROUNDS: u32 = 64;
20const SPREAD: f32 = 0.04;
21const OUTLIER: f32 = 4.0;
22const FLOOR: f32 = 0.1;
23const REACH: f32 = 0.8;
24const RANGE: f32 = 0.15;
25const PASSES: usize = 2;
26const V1_PROBES: usize = 64;
27const BATCH_LABEL: &str = "bake probe batch";
28const SAMPLES_LABEL: &str = "bake probes";
29
30type Lobes = [[f32; 3]; 6];
31type Estimates = Vec<(Lobes, [f32; 6])>;
32const FACES: [([f32; 3], [f32; 3], [f32; 3]); 6] = [
33    ([1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]),
34    ([-1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0]),
35    ([0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, -1.0]),
36    ([0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]),
37    ([0.0, 0.0, 1.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]),
38    ([0.0, 0.0, -1.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]),
39];
40
41#[derive(Clone)]
42pub struct Anchor {
43    pub hour: f32,
44    pub sky: Sky,
45    pub sun: Sun,
46}
47
48#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
49pub struct Rounds {
50    pub enabled: u32,
51    pub active: Vec<u32>,
52}
53
54impl Rounds {
55    pub fn run(&self) -> u32 {
56        self.active.len() as u32
57    }
58
59    pub fn converged(&self) -> Vec<u32> {
60        let mut before = self.enabled;
61        self.active
62            .iter()
63            .map(|&after| {
64                let done = before.saturating_sub(after);
65                before = after;
66                done
67            })
68            .collect()
69    }
70}
71
72pub struct Baked {
73    pub grid: Grid,
74    pub rounds: Rounds,
75}
76
77pub struct Paced<'a> {
78    pub pacer: &'a mut Pacer,
79    pub between: &'a mut dyn FnMut(f64) -> Result<bool, String>,
80}
81
82impl Paced<'_> {
83    pub(crate) fn quiet<T>(work: impl FnOnce(&mut Paced<'_>) -> T) -> T {
84        let mut pacer = Pacer::default();
85        let mut turns = Turns::default();
86        let mut between = |ms: f64| Ok(turns.add(ms));
87        work(&mut Paced {
88            pacer: &mut pacer,
89            between: &mut between,
90        })
91    }
92
93    pub(crate) fn run_gpu(
94        &mut self,
95        label: &str,
96        gpu: &Gpu,
97        count: u32,
98        mut encode: impl FnMut(&mut wgpu::CommandEncoder, u32, u32),
99    ) -> Result<Stats, String> {
100        if count == 0 {
101            return Ok(Stats::default());
102        }
103        let mut failed = None::<String>;
104        let between = &mut *self.between;
105        let stats = self.pacer.run(
106            label,
107            &gpu.device,
108            &gpu.queue,
109            Work::Dispatches { count },
110            |encoder, slice| {
111                if let Slice::Dispatches { start, count } = slice {
112                    encode(encoder, start, count);
113                }
114            },
115            |ms| match between(ms) {
116                Ok(turned) => turned,
117                Err(error) => {
118                    failed.get_or_insert(error);
119                    false
120                }
121            },
122        )?;
123        failed.map_or(Ok(stats), Err)
124    }
125}
126
127pub struct BakeScene {
128    pub triangles: Vec<Triangle>,
129    pub shapes: Vec<Shape>,
130    pub materials: Vec<Material>,
131    pub anchors: Vec<Anchor>,
132}
133
134fn direction(face: usize, x: usize, y: usize) -> ([f32; 3], f32) {
135    let (forward, right, up) = FACES[face];
136    let u = ((x as f32 + 0.5) / FACE as f32 * 2.0) - 1.0;
137    let v = 1.0 - (y as f32 + 0.5) / FACE as f32 * 2.0;
138    let squared = 1.0 + u * u + v * v;
139    let dir =
140        [0, 1, 2].map(|axis| (forward[axis] + right[axis] * u + up[axis] * v) / squared.sqrt());
141    let solid_angle = (4.0 / (FACE * FACE) as f32) / squared.powf(1.5) * COSINE_NORMAL;
142    (dir, solid_angle)
143}
144
145struct ProbeMeasure {
146    mean: [f32; 6],
147    second: [f32; 6],
148    backfaces: u16,
149    nearest_backface: Option<(f32, [f32; 3])>,
150}
151
152fn nearest_hit(scene: &BakeScene, bvh: &Bvh, ray: Ray) -> Option<Hit> {
153    let mesh = bvh.intersect(&scene.triangles, ray, 65504.0);
154    intersect(&scene.shapes, ray, mesh.map_or(65504.0, |hit| hit.distance)).or(mesh)
155}
156
157fn oriented_hit(scene: &BakeScene, bvh: &Bvh, sides: &[i8], ray: Ray) -> Option<(Hit, f32)> {
158    let mesh = bvh.intersect(&scene.triangles, ray, 65504.0);
159    match intersect(&scene.shapes, ray, mesh.map_or(65504.0, |hit| hit.distance)) {
160        Some(hit) => Some((hit, 1.0)),
161        None => mesh.map(|hit| (hit, f32::from(sides[hit.index as usize]))),
162    }
163}
164
165fn root(parent: &mut [u32], mut index: u32) -> u32 {
166    while parent[index as usize] != index {
167        let next = parent[parent[index as usize] as usize];
168        parent[index as usize] = next;
169        index = next;
170    }
171    index
172}
173
174fn closed_sides(triangles: &[Triangle]) -> Vec<i8> {
175    let key = |v: [f32; 3]| v.map(|value| (f64::from(value) * WELD).round() as i64);
176    let mut points: Vec<[i64; 3]> = triangles
177        .iter()
178        .flat_map(|triangle| triangle.vertices.map(key))
179        .collect();
180    points.sort_unstable();
181    points.dedup();
182    let corners: Vec<[u32; 3]> = triangles
183        .iter()
184        .map(|triangle| {
185            triangle
186                .vertices
187                .map(|v| points.binary_search(&key(v)).unwrap() as u32)
188        })
189        .collect();
190    let mut edges: Vec<(u32, u32, u32)> = corners
191        .iter()
192        .enumerate()
193        .flat_map(|(index, c)| (0..3).map(move |k| (c[k], c[(k + 1) % 3], index as u32)))
194        .collect();
195    edges.sort_unstable();
196    let span = |a: u32, b: u32| {
197        let start = edges.partition_point(|edge| (edge.0, edge.1) < (a, b));
198        let end = edges.partition_point(|edge| (edge.0, edge.1) <= (a, b));
199        &edges[start..end]
200    };
201    let mut parent: Vec<u32> = (0..triangles.len() as u32).collect();
202    let mut closed = vec![true; triangles.len()];
203    for (index, c) in corners.iter().enumerate() {
204        if c[0] == c[1] || c[1] == c[2] || c[2] == c[0] {
205            closed[index] = false;
206            continue;
207        }
208        for k in 0..3 {
209            let (a, b) = (c[k], c[(k + 1) % 3]);
210            let back = span(b, a);
211            if span(a, b).len() != 1 || back.len() != 1 {
212                closed[index] = false;
213                continue;
214            }
215            let (x, y) = (
216                root(&mut parent, index as u32),
217                root(&mut parent, back[0].2),
218            );
219            if x != y {
220                parent[x.max(y) as usize] = x.min(y);
221            }
222        }
223    }
224    let mut whole = vec![true; triangles.len()];
225    let mut volume = vec![0.0f64; triangles.len()];
226    for (index, triangle) in triangles.iter().enumerate() {
227        let group = root(&mut parent, index as u32) as usize;
228        whole[group] &= closed[index];
229        let [a, b, c] = triangle.vertices.map(|v| v.map(f64::from));
230        volume[group] += a[0] * (b[1] * c[2] - b[2] * c[1]) - a[1] * (b[0] * c[2] - b[2] * c[0])
231            + a[2] * (b[0] * c[1] - b[1] * c[0]);
232    }
233    (0..triangles.len())
234        .map(|index| {
235            let group = root(&mut parent, index as u32) as usize;
236            if !whole[group] || volume[group].abs() < 1e-15 {
237                0
238            } else if volume[group] > 0.0 {
239                1
240            } else {
241                -1
242            }
243        })
244        .collect()
245}
246
247fn measure_probe(
248    scene: &BakeScene,
249    bvh: &Bvh,
250    sides: &[i8],
251    origin: [f32; 3],
252    spacing: f32,
253) -> ProbeMeasure {
254    let mut measure = ProbeMeasure {
255        mean: [0.0; 6],
256        second: [0.0; 6],
257        backfaces: 0,
258        nearest_backface: None,
259    };
260    let limit = spacing * 4.0;
261    for face in 0..6 {
262        for y in [1, 2] {
263            for x in [1, 2] {
264                let ray = Ray {
265                    origin,
266                    direction: direction(face, x, y).0,
267                };
268                let hit = oriented_hit(scene, bvh, sides, ray);
269                let distance = hit.map_or(limit, |(hit, _)| hit.distance.min(limit));
270                measure.mean[face] += distance * 0.25;
271                measure.second[face] += distance * distance * 0.25;
272                if let Some((hit, side)) = hit
273                    && side
274                        * (0..3)
275                            .map(|axis| hit.normal[axis] * ray.direction[axis])
276                            .sum::<f32>()
277                        > 0.01
278                {
279                    measure.backfaces += 1;
280                    if measure
281                        .nearest_backface
282                        .is_none_or(|(old, _)| hit.distance < old)
283                    {
284                        measure.nearest_backface = Some((hit.distance, ray.direction));
285                    }
286                }
287            }
288        }
289    }
290    measure
291}
292
293fn probe_geometry(
294    scene: &BakeScene,
295    bvh: &Bvh,
296    sides: &[i8],
297    position: [f32; 3],
298    spacing: f32,
299) -> (ProbeExtra, [f32; 6]) {
300    let mut measure = measure_probe(scene, bvh, sides, position, spacing);
301    let mut offset = [0.0; 3];
302    if measure.backfaces > INSIDE
303        && let Some((distance, direction)) = measure.nearest_backface
304        && distance < spacing * 0.5
305    {
306        let reach = (distance + spacing * 0.1).min(spacing * 0.5);
307        offset = direction.map(|value| value * reach);
308        let relocated = std::array::from_fn(|axis| position[axis] + offset[axis]);
309        let next = measure_probe(scene, bvh, sides, relocated, spacing);
310        if next.backfaces < measure.backfaces {
311            measure = next;
312        } else {
313            offset = [0.0; 3];
314        }
315    }
316    let extra = ProbeExtra {
317        second: measure.second,
318        offset,
319        enabled: measure.backfaces <= INSIDE,
320        backfaces: measure.backfaces,
321    };
322    (extra, measure.mean)
323}
324
325pub fn bake(
326    gpu: &Gpu,
327    scene: &BakeScene,
328    spec: GridSpec,
329    samples: u32,
330    seed: u32,
331) -> Result<Vec<Grid>, String> {
332    (0..scene.anchors.len())
333        .map(|index| bake_anchor(gpu, scene, spec, samples, seed, index))
334        .collect()
335}
336
337pub fn bake_v1(
338    gpu: &Gpu,
339    scene: &BakeScene,
340    spec: GridSpec,
341    samples: u32,
342    seed: u32,
343) -> Result<Vec<Grid>, String> {
344    (0..scene.anchors.len())
345        .map(|index| bake_anchor_v1(gpu, scene, spec, samples, seed, index))
346        .collect()
347}
348
349fn bake_anchor_v1(
350    gpu: &Gpu,
351    scene: &BakeScene,
352    spec: GridSpec,
353    samples: u32,
354    seed: u32,
355    anchor_index: usize,
356) -> Result<Grid, String> {
357    if samples == 0 {
358        return Err("bake needs at least one sample per face texel".into());
359    }
360    let dims = spec.dimensions()?;
361    let count = dims.iter().map(|&v| v as usize).product::<usize>();
362    let bvh = Bvh::build(&scene.triangles);
363    let batch = Batch::new(gpu, scene, anchor_index)?;
364    let mut probes = Vec::with_capacity(count);
365    let mut turns = Turns::default();
366    for start in (0..count).step_by(V1_PROBES) {
367        let end = (start + V1_PROBES).min(count);
368        let mut records = Vec::with_capacity((end - start) * FACES.len());
369        for index in start..end {
370            let position = spec.position(dims, index);
371            let mut probe = Probe::default();
372            for (lobe, &direction) in AXES.iter().enumerate() {
373                let ray = Ray {
374                    origin: position,
375                    direction,
376                };
377                probe.visibility[lobe] =
378                    nearest_hit(scene, &bvh, ray).map_or(65504.0, |hit| hit.distance);
379            }
380            probes.push(probe);
381            for (face, &basis) in FACES.iter().enumerate() {
382                records.push(FaceRecord::new(
383                    position,
384                    basis,
385                    ray_seed(seed, anchor_index, index, face),
386                ));
387            }
388        }
389        let image = turns.time(|| batch.sample(gpu, &records, samples))?;
390        for (index, probe) in probes.iter_mut().enumerate().take(end).skip(start) {
391            for face in 0..FACES.len() {
392                let face_index = (index - start) * FACES.len() + face;
393                for y in 0..FACE {
394                    for x in 0..FACE {
395                        let tile_x = face_index % 16;
396                        let tile_y = face_index / 16;
397                        let offset = ((tile_y * FACE + y) * 64 + tile_x * FACE + x) * 16;
398                        accumulate(&mut probe.lobes, &image, offset, face, x, y);
399                    }
400                }
401            }
402        }
403    }
404    Grid::new(spec, probes)
405}
406
407pub fn bake_scaled(
408    gpu: &Gpu,
409    scene: &BakeScene,
410    spec: GridSpec,
411    samples: u32,
412    seed: u32,
413    scales: &[f32],
414) -> Result<Vec<Grid>, String> {
415    if scales.len() != scene.anchors.len() {
416        return Err("bake needs one emission scale per anchor".into());
417    }
418    scales
419        .iter()
420        .enumerate()
421        .map(|(index, &scale)| {
422            bake_anchor(
423                gpu,
424                &with_emission(scene, scale)?,
425                spec,
426                samples,
427                seed,
428                index,
429            )
430        })
431        .collect()
432}
433
434pub fn with_emission(scene: &BakeScene, scale: f32) -> Result<BakeScene, String> {
435    if !scale.is_finite() || scale < 0.0 {
436        return Err("emission scale must be finite and nonnegative".into());
437    }
438    Ok(BakeScene {
439        triangles: scene.triangles.clone(),
440        shapes: scene.shapes.clone(),
441        materials: scene
442            .materials
443            .iter()
444            .map(|material| Material {
445                emission: material.emission.map(|value| value * scale),
446                ..*material
447            })
448            .collect(),
449        anchors: scene.anchors.clone(),
450    })
451}
452
453pub fn bake_emitters(
454    gpu: &Gpu,
455    scene: &BakeScene,
456    spec: GridSpec,
457    samples: u32,
458    seed: u32,
459    direct: bool,
460) -> Result<Grid, String> {
461    Paced::quiet(|paced| bake_emitters_paced(gpu, scene, spec, samples, seed, direct, paced))
462}
463
464pub fn bake_emitters_paced(
465    gpu: &Gpu,
466    scene: &BakeScene,
467    spec: GridSpec,
468    samples: u32,
469    seed: u32,
470    direct: bool,
471    paced: &mut Paced<'_>,
472) -> Result<Grid, String> {
473    let hour = scene.anchors.first().map_or(12.0, |anchor| anchor.hour);
474    let dark = BakeScene {
475        triangles: scene.triangles.clone(),
476        shapes: scene.shapes.clone(),
477        materials: scene.materials.clone(),
478        anchors: vec![Anchor {
479            hour,
480            sky: Sky {
481                width: 1,
482                height: 1,
483                texels: vec![[0.0; 4]],
484            },
485            sun: Sun {
486                direction: [0.0, 1.0, 0.0],
487                color: [0.0; 3],
488                intensity: 0.0,
489            },
490        }],
491    };
492    bake_layer(gpu, &dark, spec, samples, seed, 0, direct, paced).map(|baked| baked.grid)
493}
494
495pub fn bake_anchor(
496    gpu: &Gpu,
497    scene: &BakeScene,
498    spec: GridSpec,
499    samples: u32,
500    seed: u32,
501    anchor_index: usize,
502) -> Result<Grid, String> {
503    bake_anchor_rounds(gpu, scene, spec, samples, seed, anchor_index).map(|baked| baked.grid)
504}
505
506pub fn bake_anchor_rounds(
507    gpu: &Gpu,
508    scene: &BakeScene,
509    spec: GridSpec,
510    samples: u32,
511    seed: u32,
512    anchor_index: usize,
513) -> Result<Baked, String> {
514    Paced::quiet(|paced| bake_anchor_paced(gpu, scene, spec, samples, seed, anchor_index, paced))
515}
516
517pub fn bake_anchor_paced(
518    gpu: &Gpu,
519    scene: &BakeScene,
520    spec: GridSpec,
521    samples: u32,
522    seed: u32,
523    anchor_index: usize,
524    paced: &mut Paced<'_>,
525) -> Result<Baked, String> {
526    bake_layer(gpu, scene, spec, samples, seed, anchor_index, true, paced)
527}
528
529fn check(scene: &BakeScene, samples: u32) -> Result<(), String> {
530    if samples == 0 {
531        return Err("bake needs at least one sample per face texel".into());
532    }
533    if scene.anchors.is_empty()
534        || scene
535            .anchors
536            .iter()
537            .any(|a| !a.hour.is_finite() || !(0.0..=24.0).contains(&a.hour))
538        || scene.anchors.windows(2).any(|a| a[0].hour >= a[1].hour)
539    {
540        return Err("bake anchors must be sorted and distinct".into());
541    }
542    Ok(())
543}
544
545fn geometry(
546    scene: &BakeScene,
547    spec: GridSpec,
548    mut between_probes: impl FnMut() -> Result<(), String>,
549) -> Result<(Vec<Probe>, Vec<ProbeExtra>), String> {
550    let dims = spec.dimensions()?;
551    let count = dims.iter().map(|&v| v as usize).product::<usize>();
552    let bvh = Bvh::build(&scene.triangles);
553    let sides = closed_sides(&scene.triangles);
554    let mut probes = Vec::with_capacity(count);
555    let mut extras = Vec::with_capacity(count);
556    for index in 0..count {
557        between_probes()?;
558        let (extra, mean) = probe_geometry(
559            scene,
560            &bvh,
561            &sides,
562            spec.position(dims, index),
563            spec.spacing,
564        );
565        probes.push(Probe {
566            visibility: mean,
567            ..Probe::default()
568        });
569        extras.push(extra);
570    }
571    Ok((probes, extras))
572}
573
574fn relocated(spec: GridSpec, dims: [u32; 3], index: usize, extra: &ProbeExtra) -> [f32; 3] {
575    let position = spec.position(dims, index);
576    std::array::from_fn(|axis| position[axis] + extra.offset[axis])
577}
578
579fn round_seed(seed: u32, round: u32) -> u32 {
580    seed.wrapping_add(round.wrapping_mul(0x27d4_eb2f))
581}
582
583fn luminance(rgb: [f32; 3]) -> f32 {
584    rgb[0] + rgb[1] + rgb[2]
585}
586
587fn kept(rounds: &[Lobes]) -> Vec<Lobes> {
588    let sums: [f32; 6] =
589        std::array::from_fn(|lobe| rounds.iter().map(|round| luminance(round[lobe])).sum());
590    let others = (rounds.len() - 1).max(1) as f32;
591    rounds
592        .iter()
593        .map(|round| {
594            let rest: [f32; 6] =
595                std::array::from_fn(|lobe| (sums[lobe] - luminance(round[lobe])) / others);
596            let floor = rest.iter().fold(0.0f32, |a, &b| a.max(b)) * FLOOR + 1e-6;
597            std::array::from_fn(|lobe| {
598                let cap = rest[lobe] * OUTLIER + floor;
599                let value = luminance(round[lobe]);
600                if value > cap {
601                    round[lobe].map(|channel| channel * cap / value)
602                } else {
603                    round[lobe]
604                }
605            })
606        })
607        .collect()
608}
609
610fn mean(rounds: &[Lobes]) -> Lobes {
611    let share = 1.0 / rounds.len().max(1) as f32;
612    std::array::from_fn(|lobe| {
613        std::array::from_fn(|channel| {
614            rounds.iter().map(|round| round[lobe][channel]).sum::<f32>() * share
615        })
616    })
617}
618
619fn estimate(rounds: &[Lobes]) -> (Lobes, [f32; 6]) {
620    if rounds.is_empty() {
621        return ([[0.0; 3]; 6], [0.0; 6]);
622    }
623    let kept = kept(rounds);
624    let count = kept.len() as f32;
625    let average = mean(&kept);
626    let error = std::array::from_fn(|lobe| {
627        let centre = luminance(average[lobe]);
628        let variance = kept
629            .iter()
630            .map(|round| (luminance(round[lobe]) - centre).powi(2))
631            .sum::<f32>()
632            / (count - 1.0).max(1.0);
633        (variance / count).sqrt()
634    });
635    (average, error)
636}
637
638fn converged(rounds: &[Lobes]) -> bool {
639    let (average, error) = estimate(rounds);
640    let level = average.map(luminance);
641    let floor = level.iter().fold(0.0f32, |a, &b| a.max(b)) * FLOOR;
642    (0..6).all(|lobe| error[lobe] <= SPREAD * level[lobe].max(floor))
643}
644
645fn converge(
646    enabled: &[bool],
647    mut sample: impl FnMut(u32, &[usize]) -> Result<Vec<Lobes>, String>,
648) -> Result<(Estimates, Rounds), String> {
649    let mut rounds: Vec<Vec<Lobes>> = vec![Vec::new(); enabled.len()];
650    let mut active: Vec<usize> = (0..enabled.len()).filter(|&i| enabled[i]).collect();
651    let mut report = Rounds {
652        enabled: active.len() as u32,
653        active: Vec::new(),
654    };
655    let mut round = 0;
656    while !active.is_empty() && round < MAX_ROUNDS {
657        let fresh = sample(round, &active)?;
658        for (&index, lobes) in active.iter().zip(fresh) {
659            rounds[index].push(lobes);
660        }
661        round += 1;
662        if round >= ROUNDS {
663            active.retain(|&index| !converged(&rounds[index]));
664        }
665        report.active.push(active.len() as u32);
666    }
667    Ok((rounds.iter().map(|list| estimate(list)).collect(), report))
668}
669
670fn smooth(
671    dims: [u32; 3],
672    spacing: f32,
673    probes: &mut [Probe],
674    extras: &[ProbeExtra],
675    errors: &[[f32; 6]],
676) {
677    for _ in 0..PASSES {
678        let source: Vec<Lobes> = probes.iter().map(|probe| probe.lobes).collect();
679        for index in 0..probes.len() {
680            if !extras[index].enabled {
681                continue;
682            }
683            let xyz = [
684                index as u32 % dims[0],
685                index as u32 / dims[0] % dims[1],
686                index as u32 / dims[0] / dims[1],
687            ];
688            let mut neighbours = Vec::with_capacity(6);
689            for axis in 0..3 {
690                for sign in [-1i32, 1] {
691                    let step = xyz[axis] as i32 + sign;
692                    if step < 0 || step >= dims[axis] as i32 {
693                        continue;
694                    }
695                    let mut other = xyz;
696                    other[axis] = step as u32;
697                    let other = GridSpec::index(dims, other);
698                    let outgoing = axis * 2 + usize::from(sign < 0);
699                    let incoming = axis * 2 + usize::from(sign > 0);
700                    if extras[other].enabled
701                        && probes[index].visibility[outgoing] >= spacing * REACH
702                        && probes[other].visibility[incoming] >= spacing * REACH
703                    {
704                        neighbours.push(other);
705                    }
706                }
707            }
708            let level = source[index]
709                .iter()
710                .map(|&rgb| luminance(rgb))
711                .fold(0.0f32, f32::max)
712                * FLOOR;
713            for (lobe, &centre) in source[index].iter().enumerate() {
714                let own = luminance(centre);
715                let mut sum = centre;
716                let mut total = 1.0;
717                for &other in &neighbours {
718                    let value = luminance(source[other][lobe]);
719                    let scale = (own.max(value).max(level).max(1e-6) * RANGE).max(
720                        2.0 * (errors[index][lobe].powi(2) + errors[other][lobe].powi(2)).sqrt(),
721                    );
722                    let weight = (-((value - own) / scale).powi(2)).exp();
723                    for (channel, sum) in sum.iter_mut().enumerate() {
724                        *sum += source[other][lobe][channel] * weight;
725                    }
726                    total += weight;
727                }
728                probes[index].lobes[lobe] = sum.map(|value| value / total);
729            }
730        }
731    }
732}
733
734#[allow(clippy::too_many_arguments)]
735fn bake_layer(
736    gpu: &Gpu,
737    scene: &BakeScene,
738    spec: GridSpec,
739    samples: u32,
740    seed: u32,
741    anchor_index: usize,
742    direct: bool,
743    paced: &mut Paced<'_>,
744) -> Result<Baked, String> {
745    check(scene, samples)?;
746    let dims = spec.dimensions()?;
747    let (mut probes, extras) = geometry(scene, spec, || (paced.between)(f64::NAN).map(|_| ()))?;
748    let mut batch = Batch::new(gpu, scene, anchor_index)?;
749    if !direct {
750        batch = batch.without_direct_emission();
751    }
752    let per_round = samples.div_ceil(ROUNDS);
753    let enabled: Vec<bool> = extras.iter().map(|extra| extra.enabled).collect();
754    let most = FACES_PER_BATCH / FACES.len();
755    let mut full_cost: Option<f64> = None;
756    let (estimates, rounds) = converge(&enabled, |round, active| {
757        let mut out = Vec::with_capacity(active.len());
758        let mut at = 0;
759        while at < active.len() {
760            let probes = paced
761                .pacer
762                .learned_units(BATCH_LABEL)
763                .map_or(V1_PROBES, |units| units as usize)
764                .clamp(1, most);
765            let chunk = &active[at..(at + probes).min(active.len())];
766            at += chunk.len();
767            let records = face_records(
768                chunk,
769                round,
770                |index| relocated(spec, dims, index, &extras[index]),
771                seed,
772                anchor_index,
773            );
774            batch.load(gpu, &records)?;
775            let target = paced.pacer.target_ms();
776            paced
777                .pacer
778                .teach(SAMPLES_LABEL, first_samples(full_cost, target, per_round));
779            let stats = paced.run_gpu(SAMPLES_LABEL, gpu, per_round, |encoder, start, count| {
780                batch.encode(gpu, encoder, records.len(), start, count);
781            })?;
782            let per_sample = stats.milliseconds.iter().sum::<f64>() / f64::from(per_round);
783            let next = batch_probes(chunk.len(), probes, per_sample, target, most);
784            if chunk.len() == probes {
785                full_cost = Some(per_sample * next as f64 / probes as f64);
786            }
787            paced.pacer.teach(BATCH_LABEL, next as u32);
788            let image = batch.read(gpu)?;
789            out.extend((0..chunk.len()).map(|slot| lobes_at(&image, slot)));
790        }
791        Ok(out)
792    })?;
793    let errors: Vec<[f32; 6]> = estimates.iter().map(|(_, error)| *error).collect();
794    for (probe, (lobes, _)) in probes.iter_mut().zip(estimates) {
795        probe.lobes = lobes;
796    }
797    smooth(dims, spec.spacing, &mut probes, &extras, &errors);
798    Ok(Baked {
799        grid: Grid::with_extras(spec, probes, extras)?,
800        rounds,
801    })
802}
803
804fn batch_probes(
805    used: usize,
806    asked: usize,
807    per_sample_ms: f64,
808    target_ms: f64,
809    most: usize,
810) -> usize {
811    if !per_sample_ms.is_finite() || per_sample_ms <= 0.0 {
812        return asked.clamp(1, most);
813    }
814    if per_sample_ms > target_ms * 1.15 {
815        ((used as f64 * target_ms / per_sample_ms).floor() as usize).clamp(1, asked.max(1))
816    } else if per_sample_ms < target_ms / 2.0 && used == asked {
817        (asked * 2).min(most)
818    } else {
819        asked.clamp(1, most)
820    }
821}
822
823fn first_samples(full_cost: Option<f64>, target_ms: f64, per_round: u32) -> u32 {
824    full_cost
825        .filter(|ms| ms.is_finite() && *ms > 0.0)
826        .map_or(1, |ms| (target_ms / ms).floor() as u32)
827        .clamp(1, per_round.max(1))
828}
829
830fn face_records(
831    chunk: &[usize],
832    round: u32,
833    origin: impl Fn(usize) -> [f32; 3],
834    seed: u32,
835    anchor_index: usize,
836) -> Vec<FaceRecord> {
837    let mut records = Vec::with_capacity(chunk.len() * FACES.len());
838    for &index in chunk {
839        let origin = origin(index);
840        for (face, &basis) in FACES.iter().enumerate() {
841            records.push(FaceRecord::new(
842                origin,
843                basis,
844                ray_seed(round_seed(seed, round), anchor_index, index, face),
845            ));
846        }
847    }
848    records
849}
850
851fn lobes_at(image: &[u8], slot: usize) -> Lobes {
852    let mut lobes = [[0.0; 3]; 6];
853    for face in 0..FACES.len() {
854        let face_index = slot * FACES.len() + face;
855        for y in 0..FACE {
856            for x in 0..FACE {
857                let tile_x = face_index % 16;
858                let tile_y = face_index / 16;
859                let offset = ((tile_y * FACE + y) * 64 + tile_x * FACE + x) * 16;
860                accumulate(&mut lobes, image, offset, face, x, y);
861            }
862        }
863    }
864    lobes
865}
866
867fn ray_seed(seed: u32, anchor_index: usize, index: usize, face: usize) -> u32 {
868    seed.wrapping_add((anchor_index as u32).wrapping_mul(0x9e37_79b9))
869        .wrapping_add((index as u32).wrapping_mul(0x85eb_ca6b))
870        .wrapping_add(face as u32)
871}
872
873fn accumulate(lobes: &mut Lobes, image: &[u8], offset: usize, face: usize, x: usize, y: usize) {
874    let rgb: [f32; 3] = [0, 1, 2].map(|channel| {
875        f32::from_le_bytes(
876            image[offset + channel * 4..offset + channel * 4 + 4]
877                .try_into()
878                .unwrap(),
879        )
880    });
881    let (dir, solid_angle) = direction(face, x, y);
882    for (lobe, axis_direction) in AXES.iter().enumerate() {
883        let cosine = (0..3)
884            .map(|axis| dir[axis] * axis_direction[axis])
885            .sum::<f32>()
886            .max(0.0);
887        for (channel, value) in rgb.iter().enumerate() {
888            lobes[lobe][channel] += value.max(0.0) * cosine * solid_angle;
889        }
890    }
891}
892
893pub fn bake_anchor_slow(
894    gpu: &Gpu,
895    scene: &BakeScene,
896    spec: GridSpec,
897    samples: u32,
898    seed: u32,
899    anchor_index: usize,
900) -> Result<Grid, String> {
901    check(scene, samples)?;
902    let anchor = scene
903        .anchors
904        .get(anchor_index)
905        .ok_or("bake anchor is missing")?;
906    let dims = spec.dimensions()?;
907    let (mut probes, extras) = geometry(scene, spec, || Ok(()))?;
908    let mut trace_scene = Scene {
909        triangles: scene.triangles.clone(),
910        shapes: scene.shapes.clone(),
911        materials: scene.materials.clone(),
912        sky: anchor.sky.clone(),
913        camera: Camera {
914            origin: [0.0; 3],
915            forward: FACES[0].0,
916            right: FACES[0].1,
917            up: FACES[0].2,
918        },
919        sun: anchor.sun,
920    };
921    let detail = Detail {
922        sampler: Sampler::Pcg,
923        ..Detail::default()
924    };
925    let mut trace = Trace::new_detailed(gpu, &trace_scene, &detail, FACE as u32, FACE as u32)?;
926    let per_round = samples.div_ceil(ROUNDS);
927    let enabled: Vec<bool> = extras.iter().map(|extra| extra.enabled).collect();
928    let (estimates, _) = converge(&enabled, |round, active| {
929        let mut out = Vec::with_capacity(active.len());
930        for &index in active {
931            let origin = relocated(spec, dims, index, &extras[index]);
932            let mut lobes = [[0.0; 3]; 6];
933            for (face, &(forward, right, up)) in FACES.iter().enumerate() {
934                trace_scene.camera = Camera {
935                    origin,
936                    forward,
937                    right,
938                    up,
939                };
940                trace.set_camera(trace_scene.camera)?;
941                trace.sample(
942                    gpu,
943                    per_round,
944                    ray_seed(round_seed(seed, round), anchor_index, index, face),
945                )?;
946                let image = trace.readback_color(gpu)?;
947                for y in 0..FACE {
948                    for x in 0..FACE {
949                        accumulate(&mut lobes, &image, (y * FACE + x) * 16, face, x, y);
950                    }
951                }
952            }
953            out.push(lobes);
954        }
955        Ok(out)
956    })?;
957    let errors: Vec<[f32; 6]> = estimates.iter().map(|(_, error)| *error).collect();
958    for (probe, (lobes, _)) in probes.iter_mut().zip(estimates) {
959        probe.lobes = lobes;
960    }
961    smooth(dims, spec.spacing, &mut probes, &extras, &errors);
962    Grid::with_extras(spec, probes, extras)
963}
964
965#[cfg(test)]
966mod tests {
967    use super::*;
968    use std::time::Instant;
969
970    fn quad(
971        triangles: &mut Vec<Triangle>,
972        a: [f32; 3],
973        b: [f32; 3],
974        c: [f32; 3],
975        d: [f32; 3],
976        material: u32,
977    ) {
978        triangles.push(Triangle {
979            vertices: [a, b, c],
980            material,
981        });
982        triangles.push(Triangle {
983            vertices: [a, c, d],
984            material,
985        });
986    }
987
988    #[test]
989    fn a_uniform_sky_integrates_to_pi_on_every_lobe() {
990        for (lobe, axis) in AXES.iter().enumerate() {
991            let mut irradiance = 0.0f64;
992            for face in 0..FACES.len() {
993                for y in 0..FACE {
994                    for x in 0..FACE {
995                        let (dir, solid_angle) = direction(face, x, y);
996                        let cosine = (0..3).map(|k| dir[k] * axis[k]).sum::<f32>().max(0.0);
997                        irradiance += f64::from(cosine * solid_angle);
998                    }
999                }
1000            }
1001            assert!(
1002                (irradiance - std::f64::consts::PI).abs() < 1e-5,
1003                "lobe {lobe} integrates a unit sky to {irradiance}"
1004            );
1005        }
1006    }
1007
1008    #[test]
1009    fn backface_majority_disables_or_relocates_inside_probes() {
1010        let scene = BakeScene {
1011            triangles: Vec::new(),
1012            shapes: vec![Shape::RoundedBox {
1013                center: [0.0; 3],
1014                half: [0.2; 3],
1015                radius: 0.0,
1016                material: 0,
1017            }],
1018            materials: vec![Material::default()],
1019            anchors: Vec::new(),
1020        };
1021        let bvh = Bvh::build(&scene.triangles);
1022        let (buried, _) = probe_geometry(&scene, &bvh, &[], [0.0; 3], 0.1);
1023        assert!(!buried.enabled);
1024        assert!(buried.backfaces > 12);
1025        let (near, _) = probe_geometry(&scene, &bvh, &[], [0.19, 0.0, 0.0], 0.1);
1026        assert!(near.enabled);
1027        assert!(near.offset.iter().map(|v| v * v).sum::<f32>().sqrt() <= 0.05);
1028        let (outside, _) = probe_geometry(&scene, &bvh, &[], [0.3, 0.0, 0.0], 0.1);
1029        assert!(outside.enabled);
1030        assert_eq!(outside.offset, [0.0; 3]);
1031    }
1032
1033    fn cube(triangles: &mut Vec<Triangle>, center: [f32; 3], half: f32, inward: bool) {
1034        let corner = |x: f32, y: f32, z: f32| {
1035            [
1036                center[0] + x * half,
1037                center[1] + y * half,
1038                center[2] + z * half,
1039            ]
1040        };
1041        let faces = [
1042            [
1043                corner(1.0, -1.0, -1.0),
1044                corner(1.0, 1.0, -1.0),
1045                corner(1.0, 1.0, 1.0),
1046                corner(1.0, -1.0, 1.0),
1047            ],
1048            [
1049                corner(-1.0, -1.0, -1.0),
1050                corner(-1.0, -1.0, 1.0),
1051                corner(-1.0, 1.0, 1.0),
1052                corner(-1.0, 1.0, -1.0),
1053            ],
1054            [
1055                corner(-1.0, 1.0, -1.0),
1056                corner(-1.0, 1.0, 1.0),
1057                corner(1.0, 1.0, 1.0),
1058                corner(1.0, 1.0, -1.0),
1059            ],
1060            [
1061                corner(-1.0, -1.0, -1.0),
1062                corner(1.0, -1.0, -1.0),
1063                corner(1.0, -1.0, 1.0),
1064                corner(-1.0, -1.0, 1.0),
1065            ],
1066            [
1067                corner(-1.0, -1.0, 1.0),
1068                corner(1.0, -1.0, 1.0),
1069                corner(1.0, 1.0, 1.0),
1070                corner(-1.0, 1.0, 1.0),
1071            ],
1072            [
1073                corner(-1.0, -1.0, -1.0),
1074                corner(-1.0, 1.0, -1.0),
1075                corner(1.0, 1.0, -1.0),
1076                corner(1.0, -1.0, -1.0),
1077            ],
1078        ];
1079        for [a, b, c, d] in faces {
1080            if inward {
1081                quad(triangles, a, d, c, b, 0);
1082            } else {
1083                quad(triangles, a, b, c, d, 0);
1084            }
1085        }
1086    }
1087
1088    #[test]
1089    fn only_closed_consistent_meshes_carry_a_side() {
1090        let mut triangles = Vec::new();
1091        cube(&mut triangles, [0.0; 3], 0.2, false);
1092        cube(&mut triangles, [1.0, 0.0, 0.0], 0.2, true);
1093        quad(
1094            &mut triangles,
1095            [-1.0, -0.5, -1.0],
1096            [-1.0, -0.5, 1.0],
1097            [2.0, -0.5, 1.0],
1098            [2.0, -0.5, -1.0],
1099            0,
1100        );
1101        let mut broken = Vec::new();
1102        cube(&mut broken, [0.0, 1.0, 0.0], 0.2, false);
1103        broken[0].vertices.swap(1, 2);
1104        triangles.extend(broken);
1105        let sides = closed_sides(&triangles);
1106        assert_eq!(&sides[..12], &[1; 12]);
1107        assert_eq!(&sides[12..24], &[-1; 12]);
1108        assert_eq!(&sides[24..], &[0; 14]);
1109        for (center, inward) in [([0.0; 3], false), ([1.0, 0.0, 0.0], true)] {
1110            let mut triangles = Vec::new();
1111            cube(&mut triangles, center, 0.2, inward);
1112            let scene = BakeScene {
1113                triangles,
1114                shapes: Vec::new(),
1115                materials: vec![Material::default()],
1116                anchors: Vec::new(),
1117            };
1118            let bvh = Bvh::build(&scene.triangles);
1119            let sides = closed_sides(&scene.triangles);
1120            let (inside, _) = probe_geometry(&scene, &bvh, &sides, center, 0.1);
1121            assert!(!inside.enabled && inside.backfaces == 24);
1122            let beside = [center[0] + 0.3, center[1], center[2]];
1123            let (outside, _) = probe_geometry(&scene, &bvh, &sides, beside, 0.1);
1124            assert!(outside.enabled && outside.backfaces == 0);
1125        }
1126    }
1127
1128    #[test]
1129    fn a_room_with_open_walls_keeps_its_probes_whatever_their_winding() {
1130        let mut triangles = Vec::new();
1131        for (a, b, c, d) in [
1132            (
1133                [-1.0, 0.0, -1.0],
1134                [1.0, 0.0, -1.0],
1135                [1.0, 0.0, 1.0],
1136                [-1.0, 0.0, 1.0],
1137            ),
1138            (
1139                [1.0, 0.0, -1.0],
1140                [1.0, 2.0, -1.0],
1141                [1.0, 2.0, 1.0],
1142                [1.0, 0.0, 1.0],
1143            ),
1144            (
1145                [-1.0, 0.0, 1.0],
1146                [1.0, 0.0, 1.0],
1147                [1.0, 2.0, 1.0],
1148                [-1.0, 2.0, 1.0],
1149            ),
1150        ] {
1151            quad(&mut triangles, a, b, c, d, 0);
1152        }
1153        let scene = BakeScene {
1154            triangles,
1155            shapes: Vec::new(),
1156            materials: vec![Material::default()],
1157            anchors: Vec::new(),
1158        };
1159        let bvh = Bvh::build(&scene.triangles);
1160        let sides = closed_sides(&scene.triangles);
1161        assert!(sides.iter().all(|&side| side == 0));
1162        let (probe, _) = probe_geometry(&scene, &bvh, &sides, [0.5, 0.1, 0.5], 0.5);
1163        assert!(probe.enabled && probe.backfaces == 0 && probe.offset == [0.0; 3]);
1164    }
1165
1166    #[test]
1167    fn rounds_cap_a_firefly_keep_sparse_light_and_stop_once_they_agree() {
1168        let mut calls = [0u32; 4];
1169        let (estimates, report) = converge(&[true, true, true, false], |round, active| {
1170            Ok(active
1171                .iter()
1172                .map(|&index| {
1173                    calls[index] += 1;
1174                    let mut lobes = [[1.0; 3]; 6];
1175                    if index == 1 && round == 1 {
1176                        lobes[2] = [500.0; 3];
1177                    }
1178                    if index == 2 {
1179                        lobes = [[if round % 2 == 0 { 0.0 } else { 2.0 }; 3]; 6];
1180                    }
1181                    lobes
1182                })
1183                .collect())
1184        })
1185        .unwrap();
1186        let lobes: Vec<Lobes> = estimates.iter().map(|(lobes, _)| *lobes).collect();
1187        assert_eq!(calls, [ROUNDS, MAX_ROUNDS, MAX_ROUNDS, 0]);
1188        assert_eq!(report.enabled, 3);
1189        assert_eq!(report.run(), MAX_ROUNDS);
1190        assert_eq!(report.active[..3], [3, 3, 3]);
1191        assert_eq!(report.active[3], 2);
1192        assert!(report.active[4..].iter().all(|&count| count == 2));
1193        let done = report.converged();
1194        assert_eq!(done.iter().sum::<u32>(), 1);
1195        assert_eq!(done[3], 1);
1196        assert_eq!(lobes[0], [[1.0; 3]; 6]);
1197        assert!(lobes[1][2].iter().all(|&v| (1.04..1.06).contains(&v)));
1198        assert_eq!(lobes[1][0], [1.0; 3]);
1199        assert_eq!(lobes[2], [[1.0; 3]; 6]);
1200        assert_eq!(lobes[3], [[0.0; 3]; 6]);
1201        assert_eq!(estimates[0].1, [0.0; 6]);
1202        assert!(estimates[2].1.iter().all(|&error| error > 0.3));
1203    }
1204
1205    #[test]
1206    fn smoothing_evens_noise_keeps_real_steps_and_stops_at_surfaces() {
1207        let spacing = 0.02;
1208        let open = Probe {
1209            visibility: [spacing * 4.0; 6],
1210            ..Probe::default()
1211        };
1212        let extra = ProbeExtra {
1213            second: [0.0; 6],
1214            offset: [0.0; 3],
1215            enabled: true,
1216            backfaces: 0,
1217        };
1218        let row = |values: [f32; 5]| -> Vec<Probe> {
1219            values
1220                .iter()
1221                .map(|&value| Probe {
1222                    lobes: [[value; 3]; 6],
1223                    ..open.clone()
1224                })
1225                .collect()
1226        };
1227        let extras = vec![extra; 5];
1228        let exact = [[0.0; 6]; 5];
1229        let mut noisy = row([1.0, 1.0, 1.1, 1.0, 1.0]);
1230        smooth([5, 1, 1], spacing, &mut noisy, &extras, &exact);
1231        assert!((noisy[2].lobes[0][0] - 1.0).abs() < 0.05);
1232        let mut rough = row([1.0, 1.0, 1.6, 1.0, 1.0]);
1233        smooth([5, 1, 1], spacing, &mut rough, &extras, &exact);
1234        assert!(rough[2].lobes[0][0] > 1.55);
1235        let mut measured = row([1.0, 1.0, 1.6, 1.0, 1.0]);
1236        smooth([5, 1, 1], spacing, &mut measured, &extras, &[[1.0; 6]; 5]);
1237        assert!(measured[2].lobes[0][0] < 1.35);
1238        let mut step = row([1.0, 1.0, 1.0, 5.0, 5.0]);
1239        smooth([5, 1, 1], spacing, &mut step, &extras, &[[0.3; 6]; 5]);
1240        assert!(step[2].lobes[0][0] < 1.001 && step[3].lobes[0][0] > 4.999);
1241        let mut wall = row([1.0, 1.0, 1.1, 1.0, 1.0]);
1242        wall[2].visibility[0] = spacing * 0.3;
1243        wall[2].visibility[1] = spacing * 0.3;
1244        smooth([5, 1, 1], spacing, &mut wall, &extras, &[[0.3; 6]; 5]);
1245        assert_eq!(wall[2].lobes[0][0], 1.1);
1246    }
1247
1248    #[test]
1249    fn batches_shrink_when_one_sample_overruns_and_grow_back_when_cheap() {
1250        assert_eq!(batch_probes(512, 512, 16.0, 3.5, 512), 112);
1251        assert_eq!(batch_probes(112, 112, 3.0, 3.5, 512), 112);
1252        assert_eq!(batch_probes(100, 100, 5.0, 3.5, 512), 70);
1253        assert_eq!(batch_probes(112, 112, 1.0, 3.5, 512), 224);
1254        assert_eq!(batch_probes(40, 112, 1.0, 3.5, 512), 112);
1255        assert_eq!(batch_probes(300, 300, 1.0, 3.5, 512), 512);
1256        assert_eq!(batch_probes(4, 4, 400.0, 3.5, 512), 1);
1257        assert_eq!(batch_probes(64, 64, f64::NAN, 3.5, 512), 64);
1258        assert_eq!(first_samples(None, 3.5, 16), 1);
1259        assert_eq!(first_samples(Some(0.5), 3.5, 16), 7);
1260        assert_eq!(first_samples(Some(0.1), 3.5, 16), 16);
1261        assert_eq!(first_samples(Some(9.0), 3.5, 16), 1);
1262        assert_eq!(first_samples(Some(f64::NAN), 3.5, 4), 1);
1263    }
1264
1265    #[test]
1266    fn sampling_shader_parses() {
1267        let module = naga::front::wgsl::parse_str(crate::PROBE_WGSL).unwrap();
1268        naga::valid::Validator::new(
1269            naga::valid::ValidationFlags::all(),
1270            naga::valid::Capabilities::all(),
1271        )
1272        .validate(&module)
1273        .unwrap();
1274    }
1275
1276    fn room() -> BakeScene {
1277        let mut triangles = Vec::new();
1278        quad(
1279            &mut triangles,
1280            [-1.5, 0.0, -1.5],
1281            [1.5, 0.0, -1.5],
1282            [1.5, 0.0, 1.5],
1283            [-1.5, 0.0, 1.5],
1284            0,
1285        );
1286        quad(
1287            &mut triangles,
1288            [-1.5, 2.0, -1.5],
1289            [0.0, 2.0, -1.5],
1290            [0.0, 2.0, 1.5],
1291            [-1.5, 2.0, 1.5],
1292            1,
1293        );
1294        quad(
1295            &mut triangles,
1296            [-1.5, 0.0, -1.5],
1297            [-1.5, 2.0, -1.5],
1298            [-1.5, 2.0, 1.5],
1299            [-1.5, 0.0, 1.5],
1300            1,
1301        );
1302        quad(
1303            &mut triangles,
1304            [1.5, 0.0, -1.5],
1305            [1.5, 2.0, -1.5],
1306            [1.5, 2.0, 1.5],
1307            [1.5, 0.0, 1.5],
1308            1,
1309        );
1310        quad(
1311            &mut triangles,
1312            [-1.5, 0.0, -1.5],
1313            [1.5, 0.0, -1.5],
1314            [1.5, 2.0, -1.5],
1315            [-1.5, 2.0, -1.5],
1316            1,
1317        );
1318        quad(
1319            &mut triangles,
1320            [-1.5, 0.0, 1.5],
1321            [1.5, 0.0, 1.5],
1322            [1.5, 2.0, 1.5],
1323            [-1.5, 2.0, 1.5],
1324            1,
1325        );
1326        BakeScene {
1327            triangles,
1328            shapes: Vec::new(),
1329            materials: vec![
1330                Material {
1331                    base: [0.9; 3],
1332                    specular: 0.0,
1333                    roughness: 1.0,
1334                    ..Material::default()
1335                },
1336                Material {
1337                    base: [0.01; 3],
1338                    specular: 0.0,
1339                    roughness: 1.0,
1340                    ..Material::default()
1341                },
1342            ],
1343            anchors: vec![Anchor {
1344                hour: 12.0,
1345                sky: Sky {
1346                    width: 4,
1347                    height: 2,
1348                    texels: vec![[0.0; 4]; 8],
1349                },
1350                sun: Sun {
1351                    direction: [0.0, 1.0, 0.0],
1352                    color: [1.0; 3],
1353                    intensity: 20.0,
1354                },
1355            }],
1356        }
1357    }
1358
1359    #[test]
1360    #[ignore = "needs a GPU; run under pgpu"]
1361    fn room_sunlight_determinism_and_timing() {
1362        let started = Instant::now();
1363        let gpu = pollster::block_on(Gpu::headless()).unwrap();
1364        let scene = room();
1365        let spec = GridSpec {
1366            min: [-0.75, 0.25, -0.75],
1367            max: [0.75, 1.75, 0.75],
1368            spacing: 0.5,
1369        };
1370        let slow_started = Instant::now();
1371        let slow = bake_anchor_slow(&gpu, &scene, spec, 4, 73, 0).unwrap();
1372        let slow_seconds = slow_started.elapsed().as_secs_f64();
1373        let batch_started = Instant::now();
1374        let first = bake(&gpu, &scene, spec, 4, 73).unwrap();
1375        let batch_seconds = batch_started.elapsed().as_secs_f64();
1376        let second = bake(&gpu, &scene, spec, 4, 73).unwrap();
1377        assert_eq!(first[0].dims, [4; 3]);
1378        assert_eq!(first[0].bytes(), slow.bytes());
1379        assert_eq!(first[0].bytes(), second[0].bytes());
1380        let dark = &first[0].probes[GridSpec::index([4; 3], [0, 0, 2])];
1381        let bright = &first[0].probes[GridSpec::index([4; 3], [3, 0, 2])];
1382        assert!(
1383            bright.lobes[3][0] > dark.lobes[3][0] * 2.0 + 0.01,
1384            "bright={} dark={}",
1385            bright.lobes[3][0],
1386            dark.lobes[3][0]
1387        );
1388        let multi_spec = GridSpec {
1389            max: [1.25, 1.75, 0.75],
1390            ..spec
1391        };
1392        let multi_slow = bake_anchor_slow(&gpu, &scene, multi_spec, 1, 73, 0).unwrap();
1393        let multi_batch = bake_anchor(&gpu, &scene, multi_spec, 1, 73, 0).unwrap();
1394        assert_eq!(multi_batch.bytes(), multi_slow.bytes());
1395        eprintln!(
1396            "bake room GPU test: slow {slow_seconds:.3} s, batch {batch_seconds:.3} s, total {:.3} s",
1397            started.elapsed().as_secs_f64(),
1398        );
1399    }
1400    fn fixed_64_bake(
1401        gpu: &Gpu,
1402        scene: &BakeScene,
1403        spec: GridSpec,
1404        samples: u32,
1405        seed: u32,
1406    ) -> Grid {
1407        let dims = spec.dimensions().unwrap();
1408        let (mut probes, extras) = geometry(scene, spec, || Ok(())).unwrap();
1409        let batch = Batch::new(gpu, scene, 0).unwrap();
1410        let mut turns = Turns::default();
1411        let enabled: Vec<bool> = extras.iter().map(|extra| extra.enabled).collect();
1412        let (estimates, _) = converge(&enabled, |round, active| {
1413            let mut out = Vec::new();
1414            for chunk in active.chunks(64) {
1415                let records = face_records(
1416                    chunk,
1417                    round,
1418                    |index| relocated(spec, dims, index, &extras[index]),
1419                    seed,
1420                    0,
1421                );
1422                let image = turns.time(|| batch.sample(gpu, &records, samples.div_ceil(ROUNDS)))?;
1423                out.extend((0..chunk.len()).map(|slot| lobes_at(&image, slot)));
1424            }
1425            Ok(out)
1426        })
1427        .unwrap();
1428        let errors: Vec<[f32; 6]> = estimates.iter().map(|(_, error)| *error).collect();
1429        for (probe, (lobes, _)) in probes.iter_mut().zip(estimates) {
1430            probe.lobes = lobes;
1431        }
1432        smooth(dims, spec.spacing, &mut probes, &extras, &errors);
1433        Grid::with_extras(spec, probes, extras).unwrap()
1434    }
1435
1436    fn paced_bake(
1437        gpu: &Gpu,
1438        scene: &BakeScene,
1439        spec: GridSpec,
1440        fixed: Option<u32>,
1441    ) -> (Grid, Vec<f64>, f64) {
1442        let mut pacer = Pacer::default();
1443        pacer.set_fixed_units(fixed);
1444        let mut slices = Vec::new();
1445        let mut turns = Turns::default();
1446        let mut between = |ms: f64| {
1447            if ms.is_finite() {
1448                slices.push(ms);
1449            }
1450            Ok(turns.add(ms))
1451        };
1452        let started = Instant::now();
1453        let grid = bake_anchor_paced(
1454            gpu,
1455            scene,
1456            spec,
1457            64,
1458            73,
1459            0,
1460            &mut Paced {
1461                pacer: &mut pacer,
1462                between: &mut between,
1463            },
1464        )
1465        .unwrap()
1466        .grid;
1467        let seconds = started.elapsed().as_secs_f64();
1468        assert!(pacer.pacing().slices <= slices.len());
1469        (grid, slices, seconds)
1470    }
1471
1472    #[test]
1473    #[ignore = "needs a GPU; run under pgpu"]
1474    fn paced_and_fixed_batches_bake_the_same_bytes() {
1475        let gpu = pollster::block_on(Gpu::headless()).unwrap();
1476        let scene = room();
1477        let spec = GridSpec {
1478            min: [-1.25, 0.25, -1.25],
1479            max: [1.25, 1.75, 1.25],
1480            spacing: 0.25,
1481        };
1482        let started = Instant::now();
1483        let fixed = fixed_64_bake(&gpu, &scene, spec, 64, 73);
1484        let fixed_seconds = started.elapsed().as_secs_f64();
1485        assert!(fixed.probes.len() > FACES_PER_BATCH / FACES.len());
1486        let reference = fixed.bytes();
1487        for samples in [1, 3, 16] {
1488            let (grid, slices, _) = paced_bake(&gpu, &scene, spec, Some(samples));
1489            assert!(grid.bytes() == reference, "{samples} samples per dispatch");
1490            assert!(!slices.is_empty());
1491        }
1492        let (paced, slices, paced_seconds) = paced_bake(&gpu, &scene, spec, None);
1493        assert!(paced.bytes() == reference, "paced dispatches");
1494        assert!(bake_anchor(&gpu, &scene, spec, 64, 73, 0).unwrap().bytes() == reference);
1495        let mut sorted = slices.clone();
1496        sorted.sort_by(f64::total_cmp);
1497        eprintln!(
1498            "probe bake, {} probes, 64 samples: 64 probes a dispatch {:.3} s; paced {} dispatches, first {:.2} ms, median {:.2} ms, longest {:.2} ms, {:.3} s",
1499            fixed.probes.len(),
1500            fixed_seconds,
1501            sorted.len(),
1502            slices[0],
1503            sorted[sorted.len() / 2],
1504            sorted[sorted.len() - 1],
1505            paced_seconds,
1506        );
1507    }
1508
1509    #[test]
1510    #[ignore = "needs a GPU; run under pgpu"]
1511    fn emission_scale_zero_removes_the_lamps_bounce() {
1512        let gpu = pollster::block_on(Gpu::headless()).unwrap();
1513        let mut triangles = Vec::new();
1514        quad(
1515            &mut triangles,
1516            [-1.0, 0.0, 1.0],
1517            [1.0, 0.0, 1.0],
1518            [1.0, 0.0, -1.0],
1519            [-1.0, 0.0, -1.0],
1520            0,
1521        );
1522        quad(
1523            &mut triangles,
1524            [-0.2, 1.0, -0.2],
1525            [0.2, 1.0, -0.2],
1526            [0.2, 1.0, 0.2],
1527            [-0.2, 1.0, 0.2],
1528            1,
1529        );
1530        let scene = BakeScene {
1531            triangles,
1532            shapes: Vec::new(),
1533            materials: vec![
1534                Material {
1535                    base: [0.8; 3],
1536                    specular: 0.0,
1537                    roughness: 1.0,
1538                    ..Material::default()
1539                },
1540                Material {
1541                    base: [0.0; 3],
1542                    emission: [5.0, 3.5, 2.25],
1543                    ..Material::default()
1544                },
1545            ],
1546            anchors: vec![Anchor {
1547                hour: 22.0,
1548                sky: Sky {
1549                    width: 1,
1550                    height: 1,
1551                    texels: vec![[0.0; 4]],
1552                },
1553                sun: Sun {
1554                    direction: [0.0, 1.0, 0.0],
1555                    color: [0.0; 3],
1556                    intensity: 0.0,
1557                },
1558            }],
1559        };
1560        let spec = GridSpec {
1561            min: [0.0, 0.2, 0.0],
1562            max: [0.0, 0.6, 0.0],
1563            spacing: 0.2,
1564        };
1565        let off = bake_scaled(&gpu, &scene, spec, 8, 11, &[0.0]).unwrap();
1566        let on = bake_scaled(&gpu, &scene, spec, 8, 11, &[1.0]).unwrap();
1567        let double = bake_scaled(&gpu, &scene, spec, 8, 11, &[2.0]).unwrap();
1568        let layer = bake_emitters(&gpu, &scene, spec, 8, 11, true).unwrap();
1569        let bounce = bake_emitters(&gpu, &scene, spec, 8, 11, false).unwrap();
1570        assert!(
1571            off[0].probes.iter().all(|probe| probe
1572                .lobes
1573                .iter()
1574                .flatten()
1575                .all(|&value| value == 0.0))
1576        );
1577        assert_eq!(layer.bytes(), on[0].bytes());
1578        for ((one, two), (direct, bounced)) in on[0]
1579            .probes
1580            .iter()
1581            .zip(&double[0].probes)
1582            .zip(layer.probes.iter().zip(&bounce.probes))
1583        {
1584            for lobe in 0..6 {
1585                for channel in 0..3 {
1586                    let value = one.lobes[lobe][channel];
1587                    assert!((two.lobes[lobe][channel] - 2.0 * value).abs() <= value * 0.01 + 1e-4);
1588                }
1589            }
1590            let up = direct.lobes[2][0];
1591            let down = direct.lobes[3][0];
1592            println!(
1593                "probe: up {up:.4} down {down:.4}, bounce only up {:.4} down {:.4}",
1594                bounced.lobes[2][0], bounced.lobes[3][0]
1595            );
1596            assert!(down > 0.0 && up > down);
1597            assert!(bounced.lobes[2][0] < up * 0.5);
1598            assert!((bounced.lobes[3][0] - down).abs() <= down * 0.05);
1599        }
1600    }
1601}