Virtual Ray Light integrator plugin for PBRT-v3

vrl.cpp

#include <iostream>
#include <memory>

#include "vrl.h"
#include "ir.h"
#include "parallel.h"
#include "scene.h"
#include "imageio.h"
#include "spectrum.h"
#include "rng.h"
#include "paramset.h"
#include "progressreporter.h"
#include "interaction.h"
#include "sampling.h"
#include "samplers/random.h"
#include "stats.h"
#include "bdpt.h"
#include "media/homogeneous.h"

namespace pbrt {
	int VrlIrRandomWalk(const Scene &scene, RayDifferential ray, Sampler &sampler,
						MemoryArena &arena, Spectrum beta, Float pdf, int maxDepth,
						TransportMode mode, Dealer<Vertex> &vertices, Dealer<Segment> &segments,
						const Vertex * prev) {
		if (maxDepth == 0) return 0;
		int bounces = 0;
		// Declare variables for forward and reverse probability densities
		Float pdfFwd = pdf, pdfRev = 0;
		while (true) {
			// Attempt to create the next subpath vertex in _path_
			MediumInteraction mi;

			VLOG(2) << "Random walk. Bounces " << bounces << ", beta " << beta <<
				", pdfFwd " << pdfFwd << ", pdfRev " << pdfRev;
			// Trace a ray and sample the medium, if any
			SurfaceInteraction isect;
			bool foundIntersection = scene.Intersect(ray, &isect);

			if (ray.medium) {
				beta *= ray.medium->Sample(ray, sampler, arena, &mi);
				Segment &segment = segments.request();
				Float rayDist = std::sqrt(Dot(ray.d, ray.d));
				segment = Segment(prev, ray.d / rayDist, ray.tMax * rayDist, foundIntersection, pdfFwd);
			}

			if (beta.IsBlack()) break;

			if (mi.IsValid()) {
				// Record medium interaction in _path_ and compute forward density
				Vertex &vertex = vertices.request();
				vertex = Vertex::CreateMedium(mi, beta, pdfFwd, *prev);
				prev = &vertex;
				if (++bounces >= maxDepth) break;

				// Sample direction and compute reverse density at preceding vertex
				Vector3f wi;
				pdfFwd = pdfRev = mi.phase->Sample_p(-ray.d, &wi, sampler.Get2D());
				ray = mi.SpawnRay(wi);
			}
			else {
				// Handle surface interaction for path generation
				if (!foundIntersection) {
					// Capture escaped rays when tracing from the camera
					if (mode == TransportMode::Radiance) {
						Vertex &vertex = vertices.request();
						vertex = Vertex::CreateLight(EndpointInteraction(ray), beta,
													 pdfFwd);
						prev = &vertex;
						++bounces;
					}
					break;
				}

				// Compute scattering functions for _mode_ and skip over medium
				// boundaries
				isect.ComputeScatteringFunctions(ray, arena, true, mode);
				if (!isect.bsdf) {
					ray = isect.SpawnRay(ray.d);
					continue;
				}

				// Initialize _vertex_ with surface intersection information
				Vertex &vertex = vertices.request();
				vertex = Vertex::CreateSurface(isect, beta, pdfFwd, *prev);
				prev = &vertex;

				if (++bounces >= maxDepth) break;

				// Sample BSDF at current vertex and compute reverse probability
				Vector3f wi, wo = isect.wo;
				BxDFType type;
				Spectrum f = isect.bsdf->Sample_f(wo, &wi, sampler.Get2D(), &pdfFwd,
												  BSDF_ALL, &type);
				VLOG(2) << "Random walk sampled dir " << wi << " f: " << f <<
					", pdfFwd: " << pdfFwd;
				if (f.IsBlack() || pdfFwd == 0.f) break;
				beta *= f * AbsDot(wi, isect.shading.n) / pdfFwd;
				VLOG(2) << "Random walk beta now " << beta;
				pdfRev = isect.bsdf->Pdf(wi, wo, BSDF_ALL);
				if (type & BSDF_SPECULAR) {
					vertex.delta = true;
					pdfRev = pdfFwd = 0;
				}
				beta *= CorrectShadingNormal(isect, wo, wi, mode);
				VLOG(2) << "Random walk beta after shading normal correction " << beta;
				ray = isect.SpawnRay(wi);
			}
		}
		return bounces;
	}

	// taken from http://geomalgorithms.com/a07-_distance.html
	Float ComputeClosestLineLine(const Point3f & p1, const Vector3f & v1, const Point3f & p2, const Vector3f & v2, Float * t1, Float * t2) {
		Float a = Dot(v1, v1);
		Float b = Dot(v1, v2);
		Float c = Dot(v2, v2);
		Vector3f w0 = p1 - p2;
		Float d = Dot(v1, w0);
		Float e = Dot(v2, w0);
		Float denominator = a * c - b * b;

		if (std::abs(denominator) <= 1e-5) {
			*t1 = d / b;
			*t2 = e / c;
		} else {
			*t1 = (b * e - c * d) / denominator;
			*t2 = (a * e - b * d) / denominator;
		}

		Point3f u = p1 + v1 * (*t1);
		Point3f v = p2 + v2 * (*t2);
		return Distance(u, v);
	}

	// taken from http://mathworld.wolfram.com/Point-LineDistance3-Dimensional.html
	Float DistanceLinePoint(const Point3f & lo, const Vector3f & ld, const Point3f & x)
	{
		Vector3f crossed = Cross(x - lo, x - lo + ld);
		return std::sqrt(Dot(crossed, crossed) / Dot(ld, ld));
	}

	Float VrlA(const Float x, const Float invH, const Float sinTheta) {
		return std::asinh(x * invH * sinTheta);
	}

	Float VrlB(const Float x, const Float h) {
		return std::atan2(x, h);
	}

	Spectrum GatherVrls(const Scene &scene,
						Sampler &sampler,
						MemoryArena &arena,
						const Camera &camera,
						const CameraSample &cameraSample,
						const Float invNumPhotonPaths,
						Dealer<Vertex> &lightVertices,
						Dealer<Segment> &lightVrls) {

		// trace camere to first sample hit
		Ray ray;
		Spectrum beta = camera.GenerateRay(cameraSample, &ray);

		Vertex c0 = Vertex::CreateCamera(&camera, ray, beta);

		Float pdfPos, pdfDir;
		camera.Pdf_We(ray, &pdfPos, &pdfDir);

		Spectrum result(0.0f);
		const int numSamplesPerVrl = 1;
		while (true) {
			// do intersection first to modify ray.tmax
			SurfaceInteraction isect;
			bool foundIntersection = scene.Intersect(ray, &isect);

			// if current ray is inside medium
			Spectrum volumeResult(0.0);
			if (std::isfinite(ray.tMax) && ray.medium) {

				if (const HomogeneousMedium * medium = dynamic_cast<const HomogeneousMedium *>(ray.medium))
				{
					Float rayDist = std::sqrt(Dot(ray.d, ray.d));
					ray.d /= rayDist;
					ray.tMax *= rayDist;

					for (unsigned int ivrl = 0; ivrl < lightVrls.size; ivrl++) {
						const Segment & vrl = lightVrls.data[ivrl];

						// compute cosTheta and sinTheta;
						Float cosTheta = Dot(vrl.d, ray.d);
						Float cosTheta2 = cosTheta * cosTheta;
						Float sinTheta2 = 1.0 - cosTheta * cosTheta;

						// compute closest points and distance
						Float uh, vh;
						Float h = ComputeClosestLineLine(ray.o, ray.d, vrl.vpl->mi.p, vrl.d, &uh, &vh);

						if (sinTheta2 <= 1e-5 || h <= 1e-5)
						{
							// special case (not in the paper) :: VRL is almost parallel with eye ray, VRL intersect eye ray

							for (int uvSamples = 0; uvSamples < numSamplesPerVrl; uvSamples++)
							{
							#if 0
								// sample VRL according to exponential distribution and treat it as VPL [0, tMax]
								int channel = std::min((int)(sampler.Get1D() * Spectrum::nSamples), Spectrum::nSamples - 1);
								Float r = sampler.Get1D();

								Float expTerm = std::exp(-medium->sigma_t[channel] * std::min(vrl.tMax, MaxFloat));
								Float vhat = -std::log(1 - r + r * expTerm) / medium->sigma_t[channel];

								Spectrum Tr = Exp(-medium->sigma_t * std::min(vhat, MaxFloat));
								Spectrum density = medium->sigma_t / (Tr - Tr / Exp(-medium->sigma_t * std::min(vrl.tMax, MaxFloat)));
								Float vPdf = 0;
								for (int i = 0; i < Spectrum::nSamples; i++) vPdf += density[i];
								vPdf *= 1 / (Float)Spectrum::nSamples;
								if (vPdf == 0)
								{
									CHECK(Tr.IsBlack());
									vPdf = 1;
								}
								Spectrum vrlBeta = Tr / vPdf;
								Float invVPdf = 1.0f / vPdf;
								Point3f v = vrl.vpl->mi.p + vhat * vrl.d;
							#else
								// uniform sampling
								Float vhat = Lerp(sampler.Get1D(), 0, vrl.tMax);
								Float invVPdf = vrl.tMax;
								Point3f v = vrl.vpl->mi.p + vhat * vrl.d;
							#endif

							#if 0
								// sample along camera ray using Kulla et al. method respect to sampled VPL
							#else
								// uniform sampling
								Float uhat = Lerp(sampler.Get1D(), 0, ray.tMax);
								Float invUPdf = ray.tMax;
								Point3f u = ray.o + uhat * ray.d;
							#endif

								Float distUv2 = DistanceSquared(u, v);
								CHECK_GT(distUv2, 0.0);
								Float distUv = std::sqrt(distUv2);
								Vector3f uv = (v - u) / distUv;

								// compute contribution
								const PhaseFunction *phase = ARENA_ALLOC(arena, HenyeyGreenstein)(medium->g);
								SimpleVisibilityTester tru(ray.o, u);
								SimpleVisibilityTester trv(vrl.vpl->p(), v);
								SimpleVisibilityTester trw(u, v);

								CHECK_GT(vrl.pdfFwd, 0.0);

								Float vplDot = 1.0f;
								if (vrl.vpl->IsLight() || vrl.vpl->IsOnSurface())
								{
									vplDot *= std::max(Dot(Vector3f(vrl.vpl->ng()), vrl.d), static_cast<Float>(0.0));
								}

								volumeResult += vrl.vpl->beta * vplDot
									* medium->sigma_s * medium->sigma_s
									* phase->p(-ray.d, uv) * phase->p(-vrl.d, -uv)
									* tru.Tr(scene, ray.time, ray.medium, sampler)
									* trv.Tr(scene, ray.time, ray.medium, sampler)
									* trw.Tr(scene, ray.time, ray.medium, sampler)
									* invUPdf * invVPdf
									/ distUv2 / vrl.pdfFwd / static_cast<Float>(numSamplesPerVrl) * beta;
							}
						}
						else
						{
							Float invH = 1.0 / h;
							Float sinTheta = std::sqrt(sinTheta2);

							// change of variables for starting and ending point 
							Float uhat0 = -uh; Float uhat1 = ray.tMax - uh;
							Float vhat0 = -vh; Float vhat1 = vrl.tMax - vh;

							// VRL importance sampling for homogeneous media

							for (int uvSamples = 0; uvSamples < numSamplesPerVrl; uvSamples++)
							{
							#if 0
								CHECK_GT(sinTheta, 0.0);

								// compute pdf normalization constant for sampling vrl
								Float avhat0 = VrlA(vhat0, invH, sinTheta);
								Float avhat1 = VrlA(vhat1, invH, sinTheta);

								// sample vhat according to inverse CDF
								Float vhat = h * std::sinh(Lerp(sampler.Get1D(), avhat0, avhat1)) / sinTheta;
								Float invVPdf = ((avhat1 - avhat0) * std::sqrt(h * h + vhat * vhat * sinTheta2)) / sinTheta;
								Point3f v = vrl.vpl->mi.p + (vhat - vhat0) * vrl.d;
							#else
								// uniform sampling
								Float vhat = Lerp(sampler.Get1D(), vhat0, vhat1);
								Float invVPdf = vhat1 - vhat0;
								Point3f v = vrl.vpl->mi.p + (vhat - vhat0) * vrl.d;
							#endif

							#if 0
								Float h_ = DistanceLinePoint(ray.o, ray.d, v);
								CHECK_GT(h_, 0.0);

								Float buhat0 = VrlB(uhat0, h_);
								Float buhat1 = VrlB(uhat1, h_);

								// sample uhat according to inverse CDF
								Float uhat = h_ * std::tan(Lerp(sampler.Get1D(), buhat0, buhat1));
								Float invUPdf = ((buhat1 - buhat0) * (h_ * h_ + uhat * uhat)) / h_;
								Point3f u = ray.o + (uhat - uhat0) * ray.d;
							#else
								// uniform sampling
								Float uhat = Lerp(sampler.Get1D(), uhat0, uhat1);
								Float invUPdf = uhat1 - uhat0;
								Point3f u = ray.o + (uhat - uhat0) * ray.d;
							#endif

								Float distUv2 = DistanceSquared(u, v);
								CHECK_GT(distUv2, 0.0);
								Float distUv = std::sqrt(distUv2);
								Vector3f uv = (v - u) / distUv;

								// compute contribution
								const PhaseFunction *phase = ARENA_ALLOC(arena, HenyeyGreenstein)(medium->g);
								SimpleVisibilityTester tru(ray.o, u);
								SimpleVisibilityTester trv(vrl.vpl->p(), v);
								SimpleVisibilityTester trw(u, v);

								CHECK_GT(vrl.pdfFwd, 0.0);

								Float vplDot = 1.0f;
								if (vrl.vpl->IsLight() || vrl.vpl->IsOnSurface())
								{
									vplDot *= std::max(Dot(Vector3f(vrl.vpl->ng()), vrl.d), static_cast<Float>(0.0));
								}
								volumeResult += vrl.vpl->beta * vplDot
									* medium->sigma_s * medium->sigma_s
									* phase->p(-ray.d, uv) * phase->p(-vrl.d, -uv)
									* tru.Tr(scene, ray.time, ray.medium, sampler)
									* trv.Tr(scene, ray.time, ray.medium, sampler)
									* trw.Tr(scene, ray.time, ray.medium, sampler)
									* invUPdf * invVPdf
									/ distUv2 / vrl.pdfFwd / static_cast<Float>(numSamplesPerVrl) * beta;
							}
						}
					}
				}
			}
			volumeResult *= invNumPhotonPaths;
			result += volumeResult;

			// check if the ray intersect any thing
			if (!foundIntersection) {
				break;
			}

			isect.ComputeScatteringFunctions(ray, arena, true);

			// similar to direct light we have to continue if
			// 1. !bsdf
			// 2. find specular surface
			if (!isect.bsdf) {
				ray = isect.SpawnRay(ray.d);
				continue;
			}

			break;
		}

		return result;

	}

	void VrlIr::Render(const Scene & scene) {
		const int photonChunkSize = 8192;

		RandomSampler sampler(nIterations);
		RandomSampler imageSampler(nIterations);

		std::unique_ptr<Distribution1D> lightDistr = ComputeLightPowerDistribution(scene);

		MemoryArena arena;

		ProgressReporter progress(nIterations, "Rendering");

		// allocate photon sampler for threads
		std::vector<std::unique_ptr<Sampler>> photonSamplers;
		for (size_t i = 0; i < photonsPerIteration; i += photonChunkSize) {
			photonSamplers.push_back(sampler.Clone(i));
		}

		Bounds2i pixelBounds = camera->film->croppedPixelBounds;
		Vector2i pixelExtent = pixelBounds.Diagonal();
		const int tileSize = 16;
		Point2i nTiles((pixelExtent.x + tileSize - 1) / tileSize, (pixelExtent.y + tileSize - 1) / tileSize);

		// initialize image
		std::unique_ptr<Spectrum[]> image(new Spectrum[pixelBounds.Area()]);
		for (size_t i = 0; i < pixelBounds.Area(); i++) { image[i] = Spectrum(0.0); }

		// allocate tile sampler for film
		std::vector<std::unique_ptr<Sampler>> tileSamplers;
		for (size_t i = 0; i < nTiles.x * nTiles.y; i++) {
			tileSamplers.push_back(imageSampler.Clone(i));
		}

		const size_t numAllVertices = photonsPerIteration * (maxDepth + 1);
		for (size_t i = 0; i < this->nIterations; i++) {
			Dealer<Vertex> lightVertices(arena.Alloc<Vertex>(numAllVertices), numAllVertices);
			Dealer<Segment> lightVrls(arena.Alloc<Segment>(numAllVertices), numAllVertices);

			// Trace photons
			ParallelFor([&](int photonIndex) {
				if (maxDepth > 0) {
					Sampler& subSampler = *photonSamplers[photonIndex / photonChunkSize];
					Float lightPdf;
					Float lightSample = subSampler.Get1D();

					int lightNum = lightDistr->SampleDiscrete(lightSample, &lightPdf);
					const std::shared_ptr<Light> &light = scene.lights[lightNum];

					Float uLightTime = Lerp(subSampler.Get1D(), camera->shutterOpen, camera->shutterClose);

					RayDifferential ray;
					Normal3f normalLight;
					Float pdfPos, pdfDir;
					Spectrum Le = light->Sample_Le(subSampler.Get2D(), subSampler.Get2D(), uLightTime, &ray, &normalLight, &pdfPos, &pdfDir);
					if (pdfPos == 0 || pdfDir == 0 || Le.IsBlack()) return;

					Spectrum beta = (AbsDot(normalLight, ray.d) * Le) / (lightPdf * pdfPos * pdfDir);
					if (beta.IsBlack()) return;

					Vertex & vertex = lightVertices.request();
					vertex = Vertex::CreateLight(light.get(), ray, normalLight, Le / (pdfPos * lightPdf), pdfPos * lightPdf);
					VrlIrRandomWalk(scene, ray, subSampler, arena, beta, pdfDir, maxDepth - 1, TransportMode::Importance, lightVertices, lightVrls, &vertex);
				}
			}, photonsPerIteration, photonChunkSize);

			// Render with photons
			{
				Float invPhotonPaths = 1.0f / static_cast<Float>(photonsPerIteration);
				ParallelFor2D([&](Point2i tile) {
					int seed = tile.y * nTiles.x + tile.x;
					Sampler& tileSampler = *tileSamplers[seed];

					int x0 = pixelBounds.pMin.x + tile.x * tileSize;
					int x1 = std::min(x0 + tileSize, pixelBounds.pMax.x);
					int y0 = pixelBounds.pMin.y + tile.y * tileSize;
					int y1 = std::min(y0 + tileSize, pixelBounds.pMax.y);
					Bounds2i tileBounds(Point2i(x0, y0), Point2i(x1, y1));

					for (Point2i pixel : tileBounds) {
						CameraSample cameraSample = tileSampler.GetCameraSample(pixel);
						Spectrum L = GatherVrls(scene, tileSampler, arena, *camera, cameraSample, invPhotonPaths, lightVertices, lightVrls);

						size_t offset = (pixelBounds.pMax.x - pixelBounds.pMin.x) * (pixel.y - pixelBounds.pMin.y) + (pixel.x - pixelBounds.pMin.x);
						image[offset] += L;
					}
				}, nTiles);
			}
			progress.Update();
			arena.Reset();
		}
		progress.Done();

		for (size_t i = 0; i < pixelBounds.Area(); i++) { image[i] /= static_cast<Float>(this->nIterations); }

		camera->film->SetImage(image.get());
		camera->film->WriteImage();
	}

	Integrator *CreateVrlIntegrator(const ParamSet &params,
									std::shared_ptr<Sampler> &sampler,
									std::shared_ptr<const Camera> &camera) {
		int np;
		int nIterations = params.FindOneInt("iterations", params.FindOneInt("numiterations", 64));
		int maxDepth = params.FindOneInt("maxdepth", 5);
		int photonsPerIter = params.FindOneInt("photonsperiteration", -1);
		int writeFreq = params.FindOneInt("imagewritefrequency", 1 << 31);

		const int *pb = params.FindInt("pixelbounds", &np);
		Bounds2i pixelBounds = camera->film->GetSampleBounds();
		if (pb) {
			if (np != 4)
				Error("Expected four values for \"pixelbounds\" parameter. Got %d.",
					  np);
			else {
				pixelBounds = Intersect(pixelBounds,
										Bounds2i{ { pb[0], pb[2] },{ pb[1], pb[3] } });
				if (pixelBounds.Area() == 0)
					Error("Degenerate \"pixelbounds\" specified.");
			}
		}
		Float rrThreshold = params.FindOneFloat("rrthreshold", 1.);
		return new VrlIr(camera, sampler, nIterations, photonsPerIter, maxDepth, writeFreq);
	}
};

vrl.hpp

#if defined(_MSC_VER)
#define NOMINMAX
#pragma once
#endif

#ifndef PBRT_INTEGRATORS_VRL_H
#define PBRT_INTEGRATORS_VRL_H

#include "pbrt.h"
#include "integrator.h"
#include "camera.h"
#include "film.h"
#include "sampler.h"
#include "bdpt.h"

namespace pbrt {

	struct Segment {
		Segment() {}

		Segment(const Vertex * vpl, const Vector3f d, const Float tMax, const bool foundIntersection, const Float pdfFwd):
			vpl(vpl), d(d), tMax(tMax), foundIntersection(foundIntersection), pdfFwd(pdfFwd)
		{
			CHECK((!std::isnan(tMax)));
		}

		const Vertex * vpl;
		Vector3f d;
		// hit surface on the other end
		bool foundIntersection;
		Float tMax;
		Float pdfFwd;
	};

	class VrlIr : public Integrator {
	public:
		VrlIr(std::shared_ptr<const Camera> & camera,
						 std::shared_ptr<Sampler> &sampler,
						 int nIterations,
						 int photonsPerIteration,
						 int maxDepth,
						 int writeFrequency) :
			camera(camera),
			sampler(sampler),
			nIterations(nIterations),
			photonsPerIteration(photonsPerIteration > 0
								? photonsPerIteration
								: camera->film->croppedPixelBounds.Area()),
			maxDepth(maxDepth),
			writeFrequency(writeFrequency) {}

		void Render(const Scene & scene);

		std::shared_ptr<Sampler> sampler;
		std::shared_ptr<const Camera> camera;
		const int nIterations;
		const int maxDepth;
		const int photonsPerIteration;
		const int writeFrequency;
	};

	Integrator *CreateVrlIntegrator(const ParamSet &params, std::shared_ptr<Sampler> &sampler, std::shared_ptr<const Camera> &camera);
}

#endif