Monte Carlo Walk-on-Spheres sampling of a vector field defined on mesh surfaces

gistfile1.txt

// Copyright Ryan Schmidt 2020
// Released under the terms of the Boost License: https://www.boost.org/LICENSE_1_0.txt

static FVector ComputeMonteCarloVectorFieldSample(
	FDynamicMesh3& Mesh, 
	FDynamicMeshAABBTree3& AABBTree, 
	FVector3d FixedDirection,	// target fixed direction. This will be projected onto mesh normals to create vector field that is interpolated
	FVector3d WorldPosition,	// position we want to sample at
	int NumSamples,				// How many monte-carlo paths to compute and average
	float Epsilon,				// path is terminated when we are within Epsilon of boundary constraints (ie mesh)
	bool bUseVertexNormals,		// whether to use vertex or face normals
	bool bProjectToZPlane,		// whether to project vectors onto Z plane (2D-ifies solution in that case)
	bool bParallel)				
{
	int32 MaxPathLen = 50;		// maximum length of Monte Carlo path before we abort (could be parameter)
	FRandomStream Random;		// for picking random points on spheres

	const FVector3d StartPos = WorldPosition;
	const FVector3d StartPosNearest = AABBTree.FindNearestPoint(StartPos);

	FCriticalSection AccumLock;
	FVector3f AccumVector = FVector3f::Zero();
	int32 AccumSamples = 0;

	ParallelFor(NumSamples, [&](int32 si)
	{
		FVector3d CurPos = StartPos;
		FVector3d NearestPos = StartPosNearest;
		IMeshSpatial::FQueryOptions QueryOptions;

		// WoS - repeatedly compute nearest point
		bool bTerminated = false;
		for (int32 pi = 0; pi < MaxPathLen && !bTerminated; ++pi)
		{
			double Dist = CurPos.Distance(NearestPos);
			if (Dist < Epsilon)
			{
				// within surface tolerance - accumulate a sample at nearest point on mesh

				double NearDistSqr;
				QueryOptions.MaxDistance = Epsilon;
				int32 tid = AABBTree.FindNearestTriangle(CurPos, NearDistSqr, QueryOptions);
				FDistPoint3Triangle3d DistQuery = TMeshQueries<FDynamicMesh3>::TriangleDistance(Mesh, tid, CurPos);

				// compute surface normal at this point
				FVector3d UseNormal = Mesh.GetTriNormal(tid);
				if (bUseVertexNormals)
				{
					FIndex3i TriVerts = Mesh.GetTriangle(tid);
					FVector3f NormalA = Mesh.GetVertexNormal(TriVerts.A);
					FVector3f NormalB = Mesh.GetVertexNormal(TriVerts.B);
					FVector3f NormalC = Mesh.GetVertexNormal(TriVerts.C);
					FVector3f VtxNormal = DistQuery.TriangleBaryCoords.X * NormalA +
						DistQuery.TriangleBaryCoords.Y * NormalB + DistQuery.TriangleBaryCoords.Z * NormalC;
					UseNormal = (FVector3d)VtxNormal.Normalized();
				}

				if (bProjectToZPlane)
				{
					UseNormal.Z = 0;
					UseNormal.Normalize();
				}

				// project fixed direction onto tangent plane perpendicular to surface normal
				FFrame3d Plane(NearestPos, UseNormal);
				FVector3d FramePt = Plane.ToPlane(Plane.Origin + FixedDirection);
				FVector3d ProjectedDir = (FramePt - Plane.Origin).Normalized();
				if (ProjectedDir.SquaredLength() < 0.9)		// handle degenerate projetion
				{
					ProjectedDir = FixedDirection;
				}

				AccumLock.Lock();
				AccumVector += (FVector3f)ProjectedDir;
				AccumSamples++;
				AccumLock.Unlock();

				bTerminated = true;
			}
			else
			{
				// jump to random point on sphere centered at current position with radius defined by closest surface point
				CurPos = CurPos + Dist * Random.GetUnitVector();
				// find new nearest pos
				NearestPos = AABBTree.FindNearestPoint(CurPos);		// (todo: we can bound this search, right? only need to search within 2*Dist?)
			}
		}
	}, (bParallel) ? EParallelForFlags::None : EParallelForFlags::ForceSingleThread );

	AccumVector /= (AccumSamples > 0) ? (float)AccumSamples : 1.0;
	return (FVector)AccumVector.Normalized();
}