libredr/light_source/
directional.rs

1use nalgebra as na;
2use pyo3::prelude::*;
3use ndarray::{s, Array4};
4use anyhow::{Result, anyhow, ensure};
5use numpy::{PyReadonlyArray1, PyArray4, IntoPyArray, PyUntypedArrayMethods};
6
7const EPS: f32 = 1e-6;
8
9/// Construct one-hot envmap with given direction and intensity
10/// # Arguments
11/// * `resolution`: a single integer
12/// * `direction`: f32 vector of 3
13/// * `intensity`: f32 scalar
14/// # Return
15/// * 3 * 6 * `resolution` * `resolution`
16pub fn directional_envmap(resolution: usize, direction: na::VectorView3<f32>, intensity: f32) -> Result<Array4<f32>> {
17  ensure!(direction.fold(true, |acc, v| { acc && v.is_finite() }),
18          "light_source::directional_envmap: `direction` is not finite");
19  let (face_id, abs_direction_max) = direction.abs().argmax();
20  let direction = direction / abs_direction_max;
21  let mut uv = match face_id {
22    0 => [direction[1], direction[2]],
23    1 => [direction[0], direction[2]],
24    2 => [direction[0], direction[1]],
25    _ => unreachable!(),
26  };
27  uv.iter_mut().for_each(|uv| {
28    *uv = ((*uv + 1.) / 2.).clamp(EPS, 1. - EPS);
29  });
30  let envmap_col = (uv[0] * resolution as f32).floor() as usize;
31  let envmap_row = resolution - 1 - (uv[1] * resolution as f32).floor() as usize;
32  let face_id = face_id * 2 + (direction[face_id] < 0.) as usize;
33  let intensity = intensity * (resolution as f32).powi(2) * direction.norm().powi(3);
34  let mut envmap = Array4::zeros((3, 6, resolution, resolution));
35  envmap.slice_mut(s![.., face_id, envmap_row, envmap_col]).fill(intensity);
36  Ok(envmap)
37}
38
39/// Python version [`directional_envmap`]
40#[pyfunction]
41#[pyo3(name = "directional_envmap")]
42pub fn py_directional_envmap<'py>(
43    py: Python<'py>,
44    resolution: usize,
45    direction: PyReadonlyArray1<f32>,
46    intensity: f32) -> Result<Bound<'py, PyArray4<f32>>> {
47  let direction = direction.try_as_matrix().ok_or(
48    anyhow!("light_source::directional_envmap: `direction` expected shape {:?}, found {:?}", [3], direction.shape()))?;
49  let envmap = directional_envmap(resolution, direction, intensity)?;
50  Ok(envmap.into_pyarray(py))
51}