@woosh/meep-engine/src/engine/graphics/sh3/path_tracer/PathTracedScene.js

Pure JavaScript game engine. Fully featured and production ready.

import { assert } from "../../../../core/assert.js";
import { BVH, NULL_NODE } from "../../../../core/bvh2/bvh3/BVH.js";
import { ebvh_build_hierarchy } from "../../../../core/bvh2/bvh3/ebvh_build_hierarchy.js";
import { ebvh_optimize_treelet } from "../../../../core/bvh2/bvh3/ebvh_optimize_treelet.js";
import { bvh_query_user_data_ray_segment } from "../../../../core/bvh2/bvh3/query/bvh_query_user_data_ray_segment.js";
import { Cache } from "../../../../core/cache/Cache.js";
import { array_copy } from "../../../../core/collection/array/array_copy.js";
import { array_quick_sort_by_lookup_map } from "../../../../core/collection/array/array_quick_sort_by_lookup_map.js";
import { strictEquals } from "../../../../core/function/strictEquals.js";
import { AABB3 } from "../../../../core/geom/3d/aabb/AABB3.js";
import { v3_morton_encode_bounded } from "../../../../core/geom/3d/morton/v3_morton_encode_bounded.js";
import { Ray3 } from "../../../../core/geom/3d/ray/Ray3.js";
import { ray3_array_compose } from "../../../../core/geom/3d/ray/ray3_array_compose.js";
import { v3_dot } from "../../../../core/geom/vec3/v3_dot.js";
import { v3_length } from "../../../../core/geom/vec3/v3_length.js";
import { saturate } from "../../../../core/math/clamp01.js";
import { max2 } from "../../../../core/math/max2.js";
import { compute_geometry_polycount } from "../../geometry/compute_geometry_polycount.js";
import { AbstractLight } from "../../render/forward_plus/model/AbstractLight.js";
import { BufferedGeometryBVH } from "./BufferedGeometryBVH.js";
import { make_vector3 } from "./make_one_vector3.js";
import { MaterialConverter } from "./material/MaterialConverter.js";
import { PathTracedMesh } from "./PathTracedMesh.js";
import { sample_material } from "./texture/sample_material.js";

function light_getDistanceAttenuation(distance, cutoff_distance, decay_exponent) {

    let distanceFalloff = 1.0 / Math.max(Math.pow(distance, decay_exponent), 0.01);

    if (cutoff_distance > 0.0) {
        const b = distance / cutoff_distance;
        const b2 = b * b;
        const b4 = b2 * b2;

        const c = saturate(1.0 - b4);

        distanceFalloff *= c * c;
    }

    return distanceFalloff;
}

/**
 *
 * @type {number[]|Uint32Array}
 */
const scratch_uint32_array = new Uint32Array(4096);

const _ray_1 = new Ray3();

const DEFAULT_CONVERTER = new MaterialConverter();

export class PathTracedScene {


    /**
     *
     * @type {BVH}
     */
    bvh_top_level = new BVH();


    /**
     *
     * @type {Map<number, PathTracedMesh>}
     */
    meshes = new Map();

    /**
     *
     * @type {AbstractLight[]}
     */
    __lights = [];

    /**
     *
     * @type {Map<THREE.BufferGeometry, BufferedGeometryBVH>}
     */
    geo_cache = new Map();

    /**
     *
     * @type {Map<PathTracedMesh, number>}
     */
    #mesh_bvh_nodes = new Map();


    /**
     *
     * @type {Cache<THREE.Material, StandardMaterial>}
     */
    #material_cache = new Cache({
        keyEqualityFunction: strictEquals,
        keyHashFunction(mat) {
            return mat.id
        }
    });
    /**
     *
     * @type {function}
     * @private
     */
    __background_sampler = make_vector3(0, 0, 0);

    /**
     *
     * @returns {PathTracedMesh[]}
     */
    get #meshes() {
        return Array.from(this.meshes.values());
    }

    #rebuild_bvh() {
        const meshes = this.#meshes;

        const bounds = new AABB3();

        const bvh = this.bvh_top_level;

        if (bvh.root === NULL_NODE) {
            return;
        }

        bvh.node_get_aabb(bvh.root, bounds);

        bvh.release_all();

        const mesh_count = meshes.length;
        const morton_codes = new Map();

        for (let i = 0; i < mesh_count; i++) {
            const mesh = meshes[i];

            const center = mesh.aabb.getCenter();

            const code = v3_morton_encode_bounded(center.x, center.y, center.z, bounds);

            morton_codes.set(mesh, code);
        }

        //
        array_quick_sort_by_lookup_map(meshes, morton_codes, 0, mesh_count - 1);

        const node_leaf_count = mesh_count;
        const node_bin_count = max2(0, node_leaf_count - 1);

        const node_total_count = node_leaf_count + node_bin_count;

        const nodes = new Uint32Array(node_total_count);

        // skip allocation calls, allocate exactly as many nodes as we need
        bvh.node_capacity = node_total_count;
        bvh.__size = node_total_count;

        for (let i = 0; i < node_total_count; i++) {
            // store nodes in reverse order so that top-level nodes end up on top
            nodes[i] = (node_total_count - 1) - i;
        }

        // assign leaves
        for (let i = 0; i < node_leaf_count; i++) {
            const node = nodes[i];

            const mesh = meshes[i];

            bvh.node_set_aabb(node, mesh.aabb);

            bvh.node_set_child1(node, NULL_NODE); //mark node as a leaf
            bvh.node_set_user_data(node, mesh.id);

            bvh.node_set_height(node, 0);

        }

        // record newly generated nodes as "unprocessed"
        const unprocessed_nodes = new Uint32Array(node_leaf_count);
        array_copy(nodes, 0, unprocessed_nodes, 0, node_leaf_count);

        // assign root
        bvh.__root = ebvh_build_hierarchy(bvh, unprocessed_nodes, node_leaf_count, nodes, node_leaf_count,16);

        ebvh_optimize_treelet(bvh);
    }

    optimize() {
        this.meshes.forEach((mesh, mesh_id) => {
            mesh.build_tight_bounds();
        });

        // re-build bvh from bottom up
        this.#rebuild_bvh();
    }

    /**
     *
     * @param {THREE.BufferGeometry} geo
     * @return {BufferedGeometryBVH}
     */
    obtainGeometryBVH(geo) {
        const cached = this.geo_cache.get(geo);

        if (cached !== undefined) {
            return cached;
        }

        const bvh = new BufferedGeometryBVH();

        const label = `bvh build ${compute_geometry_polycount(geo)}`;
        console.time(label);
        bvh.build(geo);
        console.timeEnd(label);

        this.geo_cache.set(geo, bvh);

        return bvh;
    }

    /**
     *
     * @param {AbstractLight} light
     */
    addLight(light) {
        assert.isInstanceOf(light, AbstractLight, 'light');

        this.__lights.push(light);
    }

    /**
     *
     * @param {PathTracedMesh} mesh
     * @returns {boolean}
     */
    hasMesh(mesh) {
        return this.#mesh_bvh_nodes.has(mesh);
    }

    /**
     *
     * @param {PathTracedMesh} mesh
     * @returns {boolean}
     */
    addMesh(mesh) {
        if (this.hasMesh(mesh)) {
            // already exists
            return false;
        }

        mesh.bvh = this.obtainGeometryBVH(mesh.geometry);

        mesh.update_bounds();

        const bvh = this.bvh_top_level;

        const bvh_node_id = bvh.allocate_node();

        this.#mesh_bvh_nodes.set(mesh, bvh_node_id);

        bvh.node_set_aabb(
            bvh_node_id,
            mesh.aabb
        );

        bvh.node_set_user_data(
            bvh_node_id,
            mesh.id
        );
        bvh.insert_leaf(bvh_node_id);

        this.meshes.set(mesh.id, mesh);


        return true;
    }

    /**
     *
     * @param {PathTracedMesh} mesh
     * @returns {boolean}
     */
    removeMesh(mesh) {

        if (!this.hasMesh(mesh)) {
            return false;
        }

        const node_id = this.#mesh_bvh_nodes.get(mesh);

        this.#mesh_bvh_nodes.delete(mesh);

        this.bvh_top_level.remove_leaf(node_id);
        this.bvh_top_level.release_node(node_id);

        return true;
    }

    /**
     * Retrieves pre-cached material or build one from scratch, caches it and returns the result
     * @param {THREE.Material} material
     * @returns {StandardMaterial}
     */
    obtainMaterial(material) {
        return this.#material_cache.getOrCompute(material, DEFAULT_CONVERTER.convert, DEFAULT_CONVERTER);
    }

    /**
     *
     * @param {THREE.BufferGeometry} geometry
     * @param {THREE.Material} material
     * @param {mat4|number[]} transform
     */
    createMesh(geometry, material, transform) {
        const standard_material = this.obtainMaterial(material);

        const mesh = new PathTracedMesh();
        mesh.geometry = geometry;

        mesh.material = standard_material;

        mesh.transform = transform;

        this.addMesh(mesh);
    }

    /**
     *
     * @param {number[]} out [color_r, color_g, color_b, normal_x, normal_y, normal_z]
     * @param {number[]} hit
     * @param {Ray3} incoming_ray
     */
    sample_material(out, hit, incoming_ray) {
        const primitive_id = hit[9];
        const instance_id = hit[10];
        const u = hit[7];
        const v = hit[8];

        out[0] = 1;
        out[1] = 1;
        out[2] = 1;

        /**
         *
         * @type {PathTracedMesh}
         */
        const mesh = this.meshes.get(instance_id);

        if (mesh === undefined) {
            // instance not found
            return;
        }

        sample_material(out, mesh, primitive_id, u, v);

        if (v3_dot(
            hit[3], hit[4], hit[5], //geometry nornal
            incoming_ray[3], incoming_ray[4], incoming_ray[5],
        ) > 0) {
            // back-facing triangle
            out[3] = -out[3];
            out[4] = -out[4];
            out[5] = -out[5];
        }

        // const pdf = mesh.material.scattering_pdf(
        //     incoming_ray[3],incoming_ray[4],incoming_ray[5],
        //     out[3],out[4],out[5],
        //     incoming_ray[3],incoming_ray[4],incoming_ray[5],
        // );


    }

    /**
     * Tests ray for occlusion
     * @param {Ray3} ray
     * @returns {boolean}
     */
    occluded(ray) {

        const bvh = this.bvh_top_level;

        const ray_origin_x = ray[0];
        const ray_origin_y = ray[1];
        const ray_origin_z = ray[2];

        const ray_direction_x = ray[3];
        const ray_direction_y = ray[4];
        const ray_direction_z = ray[5];

        const max_distance = ray[6];

        const hit_count = bvh_query_user_data_ray_segment(
            bvh, bvh.root,
            scratch_uint32_array, 0,
            ray_origin_x, ray_origin_y, ray_origin_z,
            ray_direction_x, ray_direction_y, ray_direction_z,
            0, max_distance
        );


        for (let i = 0; i < hit_count; i++) {

            const node_user_data = scratch_uint32_array[i];

            const mesh = this.meshes.get(node_user_data);

            if (mesh.occluded(ray)) {
                return true;
            }

        }

        return false;
    }

    /**
     *
     * @param {number[]} out
     * @param {number[]|Ray3} ray
     * @return {number} distance to contact, or -1 if no contact found
     */
    trace(out, ray) {

        const bvh = this.bvh_top_level;

        const ray_origin_x = ray[0];
        const ray_origin_y = ray[1];
        const ray_origin_z = ray[2];

        const ray_direction_x = ray[3];
        const ray_direction_y = ray[4];
        const ray_direction_z = ray[5];

        const max_distance = ray[6];

        const hit_count = bvh_query_user_data_ray_segment(
            bvh, bvh.root,
            scratch_uint32_array, 0,
            ray_origin_x, ray_origin_y, ray_origin_z,
            ray_direction_x, ray_direction_y, ray_direction_z,
            0, max_distance
        );

        let nearest_hit_distance = max_distance;

        for (let i = 0; i < hit_count; i++) {

            const node_user_data = scratch_uint32_array[i];

            const mesh = this.meshes.get(node_user_data);

            const distance_to_hit = mesh.hit(out, ray, nearest_hit_distance);

            if (distance_to_hit >= 0) {
                // since raycast in leaf nodes is already bound by maximum distance, any hit we get is necessarily a closer hit than before
                nearest_hit_distance = distance_to_hit;
            }
        }

        if (nearest_hit_distance !== max_distance) {

            return nearest_hit_distance;

        }

        return -1;
    }

    /**
     *
     * @param {number[]} out
     * @param {number} out_offset
     * @param {number[]} direction
     * @param {number} direction_offset
     */
    sample_background(out, out_offset, direction, direction_offset) {
        this.__background_sampler(out, out_offset, direction, direction_offset);
    }

    /**
     *
     * @param {number[]} out
     * @param {number} out_offset
     * @param {number[]} ray
     */
    sample_lights(out, out_offset, ray) {

        const lights = this.__lights;
        const light_count = lights.length;

        for (let i = 0; i < 3; i++) {
            // initialize contribution to 0
            out[out_offset + i] = 0;
        }

        const ray_origin_x = ray[0];
        const ray_origin_y = ray[1];
        const ray_origin_z = ray[2];

        const ray_direction_x = ray[3];
        const ray_direction_y = ray[4];
        const ray_direction_z = ray[5];

        for (let i = 0; i < light_count; i++) {
            const light = lights[i];

            if (light.isDirectionalLight === true) {
                const dir = light.direction;

                // see https://github.com/mrdoob/three.js/blob/f0a9e0cf90a2f1ba5017fcb7fd46f02748b920cf/src/renderers/shaders/ShaderChunk/lights_physical_pars_fragment.glsl.js#L172

                const light_dir_inv_x = -dir.x;
                const light_dir_inv_y = -dir.y;
                const light_dir_inv_z = -dir.z;

                const dotNL = v3_dot(
                    ray_direction_x, ray_direction_y, ray_direction_z,
                    light_dir_inv_x, light_dir_inv_y, light_dir_inv_z
                );

                if (dotNL <= 0) {
                    // no contribution, facing away from the light
                    continue;
                }


                ray3_array_compose(_ray_1,
                    ray_origin_x, ray_origin_y, ray_origin_z,
                    light_dir_inv_x, light_dir_inv_y, light_dir_inv_z
                );

                _ray_1.tMax = Infinity;

                // check if there are any obstacles in the direction of the light
                if (this.occluded(_ray_1)) {
                    // light is occluded
                    continue;
                }


                const intensity = dotNL * light.intensity.getValue();
                const light_color = light.color;

                out[out_offset] += light_color.r * intensity;
                out[out_offset + 1] += light_color.g * intensity;
                out[out_offset + 2] += light_color.b * intensity;


            } else if (light.isPointLight === true) {

                const light_offset_x = ray_origin_x - light.position.x;
                const light_offset_y = ray_origin_y - light.position.y;
                const light_offset_z = ray_origin_z - light.position.z;

                const distance_to_light = v3_length(light_offset_x, light_offset_y, light_offset_z);

                const radius = light.radius.getValue();

                if (distance_to_light > radius) {
                    // too far
                    continue;
                }

                const light_norm_inv = 1 / distance_to_light;

                // see https://github.com/mrdoob/three.js/blob/f0a9e0cf90a2f1ba5017fcb7fd46f02748b920cf/src/renderers/shaders/ShaderChunk/lights_physical_pars_fragment.glsl.js#L172

                const light_dir_inv_x = -light_offset_x * light_norm_inv;
                const light_dir_inv_y = -light_offset_y * light_norm_inv;
                const light_dir_inv_z = -light_offset_z * light_norm_inv;

                const dotNL = v3_dot(
                    ray_direction_x, ray_direction_y, ray_direction_z,
                    light_dir_inv_x, light_dir_inv_y, light_dir_inv_z
                );

                if (dotNL <= 0) {
                    // no contribution, facing away from the light
                    continue;
                }
                ray3_array_compose(_ray_1,
                    ray_origin_x, ray_origin_y, ray_origin_z,
                    light_dir_inv_x, light_dir_inv_y, light_dir_inv_z
                );

                _ray_1.tMax = radius;

                // check if there are any obstacles in the direction of the light
                if (this.occluded(_ray_1)) {
                    // light is occluded
                    continue;
                }

                const attenuation = light_getDistanceAttenuation(distance_to_light, radius, 2);

                const intensity = dotNL * light.intensity.getValue() * attenuation;
                const light_color = light.color;

                out[out_offset] += light_color.r * intensity;
                out[out_offset + 1] += light_color.g * intensity;
                out[out_offset + 2] += light_color.b * intensity;
            }
        }
    }
}