ray-tracing2/ray-tracing-material-visualizer/src/main.rs

363 lines
11 KiB
Rust

use core::f64;
use plotters::{
chart::ChartBuilder,
coord::Shift,
prelude::{BitMapBackend, DrawingArea, DrawingBackend, IntoDrawingArea},
style::WHITE,
};
use rand::{rngs::SmallRng, SeedableRng};
use ray_tracing_core::prelude::*;
use ray_tracing_material::{
diffuse::DiffuseMaterial,
microfacet::{BeckmannDistribution, Microfacet},
mirror::Mirror,
oren_nayar::OrenNayar,
};
use rayon::iter::{
IntoParallelIterator, IntoParallelRefIterator, ParallelExtend, ParallelIterator,
};
fn main() -> Result<(), Box<dyn std::error::Error>> {
let w_in = Dir3::new(1.0, 1.0, 0.0).normalize();
let color = Color::new(1.0, 1.0, 1.0);
let m = Mirror::new(color);
generate_chart("mirror.png", &m, 2, w_in)?;
let m = DiffuseMaterial::new(color);
generate_chart("diffuse.png", &m, 100, w_in)?;
let m = Microfacet::new(BeckmannDistribution::new(0.1), color);
generate_chart("microfacet.png", &m, 100, w_in)?;
let m = OrenNayar::new(0.5, color);
generate_chart("oren-nayar.png", &m, 100, w_in)?;
Ok(())
}
fn generate_chart<M: Material<SmallRng>>(
path: impl AsRef<std::path::Path>,
m: &M,
samples_per_pixel: usize,
w_in: Dir3,
) -> Result<(), Box<dyn std::error::Error>> {
let inner_size = 400;
let width = (inner_size + 70) * 3;
let height = ((inner_size + 97) * 2 + 32) * 2;
let root = BitMapBackend::new(&path, (width, height)).into_drawing_area();
root.fill(&WHITE)?;
let executions = inner_size as usize * inner_size as usize * samples_per_pixel;
let areas = root.split_evenly((2, 1));
let eval_histogram = create_histogram_evaled(m, 400, 400, executions, w_in);
let sample_histogram = create_historgram_sampled(m, 400, 400, executions, w_in);
let max = Float::max(
Float::max(eval_histogram.max(), sample_histogram.max()),
Float::MIN_POSITIVE,
);
let eval_energy = eval_histogram.energy();
let area = areas[0].titled(
&format!(
"Evaled ({:.4}, {:.4}, {:.4})",
eval_energy.r(),
eval_energy.g(),
eval_energy.b()
),
("sans-serif", 30),
)?;
plot_material(&area, eval_histogram, max)?;
let sample_energy = sample_histogram.energy();
let area = areas[1].titled(
&format!(
"Evaled ({:.4}, {:.4}, {:.4})",
sample_energy.r(),
sample_energy.g(),
sample_energy.b()
),
("sans-serif", 30),
)?;
plot_material(&area, sample_histogram, max)?;
root.present()?;
Ok(())
}
fn plot_material<DB: DrawingBackend>(
root: &DrawingArea<DB, Shift>,
histogram: Histogram,
max: Float,
) -> Result<(), Box<dyn std::error::Error>>
where
DB::ErrorType: 'static,
{
let areas = root.split_evenly((2, 3));
for (channel, upper, lower) in (0..3).map(|i| (i, &areas[i], &areas[i + 3])) {
let channel_name = match channel {
0 => "red",
1 => "green",
2 => "blue",
_ => unreachable!(),
};
let mut upper_chart = ChartBuilder::on(upper)
.margin(5)
.x_label_area_size(30)
.top_x_label_area_size(30)
.y_label_area_size(30)
.right_y_label_area_size(30)
.caption(
format!("Upper Hemisphere {channel_name}"),
("sans-serif", 20),
)
.build_cartesian_2d(-1.0..1.0, -1.0..1.0)?;
upper_chart
.configure_mesh()
.disable_x_mesh()
.disable_y_mesh()
.draw()?;
let upper_area = upper_chart.plotting_area();
let mut lower_chart = ChartBuilder::on(lower)
.margin(5)
.x_label_area_size(30)
.top_x_label_area_size(30)
.y_label_area_size(30)
.right_y_label_area_size(30)
.caption(
format!("Lower Hemisphere {channel_name}"),
("sans-serif", 20),
)
.build_cartesian_2d(-1.0..1.0, -1.0..1.0)?;
lower_chart
.configure_mesh()
.disable_x_mesh()
.disable_y_mesh()
.draw()?;
let lower_area = lower_chart.plotting_area();
let (xrange, yrange) = upper_area.get_pixel_range();
let width = (xrange.end - xrange.start) as usize;
let height = (yrange.end - yrange.start) as usize;
assert_eq!(width, histogram.width);
assert_eq!(height, histogram.height);
let (xrange, yrange) = lower_area.get_pixel_range();
let width = (xrange.end - xrange.start) as usize;
let height = (yrange.end - yrange.start) as usize;
assert_eq!(width, histogram.width);
assert_eq!(height, histogram.height);
for (dir, f) in histogram.iter() {
let f = match channel {
0 => f.r(),
1 => f.g(),
2 => f.b(),
_ => unreachable!(),
};
assert!(f >= 0.0 && f <= max && f.is_finite() && max != 0.0, "{f}");
let color = plotters::prelude::ViridisRGB::get_color_normalized(f, 0.0, max);
if dir.y() < 0.0 {
lower_area.draw_pixel((dir.x() as f64, dir.z() as f64), &color)?;
} else {
upper_area.draw_pixel((dir.x() as f64, dir.z() as f64), &color)?;
}
}
}
Ok(())
}
struct Histogram {
data: Vec<Color>,
inserted: usize,
width: usize,
height: usize,
max: f64,
}
impl Histogram {
fn new(width: usize, height: usize) -> Self {
Self {
data: vec![Color::black(); 2 * width * height],
inserted: 0,
width,
height,
max: 0.0,
}
}
fn discretize(&self, p: Dir3) -> usize {
assert!(p.x() >= -1.0);
assert!(p.x() <= 1.0);
assert!(p.y() >= -1.0);
assert!(p.y() <= 1.0);
assert!(p.z() >= -1.0);
assert!(p.z() <= 1.0);
(((p.x() + 1.0) * 0.5 * self.width as f32).floor() as usize).min(self.width - 1)
+ (((p.z() + 1.0) * 0.5 * self.height as f32).floor() as usize).min(self.height - 1)
* self.width
+ if p.y() < 0.0 {
self.width * self.height
} else {
0
}
}
fn insert(&mut self, pos: Dir3, value: Color) {
let index = self.discretize(pos);
self.data[index] += value;
self.max = f64::max(self.max, self.data[index].r() as f64);
self.max = f64::max(self.max, self.data[index].g() as f64);
self.max = f64::max(self.max, self.data[index].b() as f64);
self.inserted += 1;
}
fn max(&self) -> f32 {
dbg!(self.max, self.inserted, self.inserted as f64);
(self.max / (self.inserted as f64)) as f32
}
fn iter(&self) -> impl Iterator<Item = (Dir3, Color)> + use<'_> {
let m = 1f32 / self.inserted as f32;
let width = self.width;
let height = self.height;
self.data
.iter()
.enumerate()
.map(move |(i, &f)| {
if i < width * height {
let x = i % width;
let z = i / width;
let x = 2.0 * (x as f32 + 0.5) / (width as f32) - 1.0;
let z = 2.0 * (z as f32 + 0.5) / (height as f32) - 1.0;
let y = f32::sqrt(1.0 - x * x - z * z);
(Dir3::new(x, y, z), f * m)
} else {
let x = i % width;
let z = (i - width * height) / width;
let x = 2.0 * (x as f32 + 0.5) / (width as f32) - 1.0;
let z = 2.0 * (z as f32 + 0.5) / (height as f32) - 1.0;
let y = -f32::sqrt(1.0 - x * x - z * z);
(Dir3::new(x, y, z), f * m)
}
})
.filter(|(dir, _)| dir.x() * dir.x() + dir.z() * dir.z() < 1.0)
}
fn energy(&self) -> Color {
let total = self.data.par_iter().copied().sum::<Color>();
total / (self.inserted as Float)
}
}
impl ParallelExtend<(Dir3, Color)> for Histogram {
fn par_extend<I>(&mut self, par_iter: I)
where
I: IntoParallelIterator<Item = (Dir3, Color)>,
{
let (tx, rx) = std::sync::mpsc::sync_channel(100);
std::thread::scope(move |s| {
s.spawn(move || {
while let Ok((i, f)) = rx.recv() {
self.insert(i, f)
}
});
{
let tx = tx;
par_iter.into_par_iter().for_each(|t| {
let _ = tx.send(t);
});
}
});
}
}
fn create_historgram_sampled<M: Material<SmallRng>>(
m: &M,
width: usize,
height: usize,
executions: usize,
w_in: Dir3,
) -> Histogram {
let mut h = Histogram::new(width, height);
h.par_extend(
(0..executions)
.into_par_iter()
.map_init(SmallRng::from_entropy, |rng, _| {
let sample = m.sample(w_in, rng);
assert!(
sample.color().r() >= 0.0
&& sample.color().r().is_finite()
&& sample.color().g() >= 0.0
&& sample.color().g().is_finite()
&& sample.color().b() >= 0.0
&& sample.color().b().is_finite()
&& (-1.0..=1.0).contains(&sample.w_out().x())
&& (-1.0..=1.0).contains(&sample.w_out().y())
&& (-1.0..=1.0).contains(&sample.w_out().z()),
"invalid_sampling: w_in: {w_in:?}; w_out: {:?}; material: {m:?}; color: {:?}",
sample.w_out(),
sample.color()
);
(sample.w_out(), sample.color())
}),
);
h
}
fn create_histogram_evaled<M: Material<SmallRng>>(
m: &M,
width: usize,
height: usize,
executions: usize,
w_in: Dir3,
) -> Histogram {
let mut h = Histogram::new(width, height);
h.par_extend(
(0..executions)
.into_par_iter()
.map_init(SmallRng::from_entropy, |rng, _| {
let w_out = Dir3::sample_uniform_sphere(rng);
let color = m.eval(w_in, w_out, rng) * w_out.y().abs() * 4.0 * FloatConsts::PI;
assert!(
color.r() >= 0.0
&& color.r().is_finite()
&& color.g() >= 0.0
&& color.g().is_finite()
&& color.b() >= 0.0
&& color.b().is_finite(),
"invalid_eval: w_in: {w_in:?}; w_out: {w_out:?}; material: {m:?}; color: {color:?}"
);
(w_out, color)
}),
);
h
}