Simplify normal estimation

Author: iwoplaza

apps/phoure-www/src/lib-ray-marching/estimateNormal.ts

@@ -1,36 +1,28 @@
 import tgpu from 'typegpu';
-import { vec3f } from 'typegpu/data';
-import { mul, normalize } from 'typegpu/std';
+import { vec2f, vec3f } from 'typegpu/data';
+import { normalize } from 'typegpu/std';
 
 import { MarchParams } from './marchSdf.ts';
-import { ShapeContext } from './types.ts';
+
+// Doing it in a derived until WGSL generation can properly decide when to `let` and when to `var`
+const epsilon = tgpu['~unstable'].derived(() =>
+  // Arbitrary - should be smaller than any surface detail in your distance function, but not so small as to get lost in float precision
+  vec2f(MarchParams.surfaceThreshold.$ * 0.5, 0),
+);
 
 /**
  * Estimates the normal vector at a point in space, based on the SDF given by `MarchParams.sampleSdf`.
  */
-export const estimateNormal = tgpu['~unstable'].fn(
-  [vec3f, ShapeContext],
+export const estimateNormal = tgpu.fn(
+  [vec3f],
   vec3f,
-)((point, ctx) => {
-  /** Arbitrary - should be smaller than any surface detail in your distance function, but not so small as to get lost in float precision */
-  const epsilon = MarchParams.getSurfaceThreshold.value(ctx) * 0.5;
-  const offX = vec3f(point.x + epsilon, point.y, point.z);
-  const offY = vec3f(point.x, point.y + epsilon, point.z);
-  const offZ = vec3f(point.x, point.y, point.z + epsilon);
-
-  const centerDistance = MarchParams.sampleSdf.value(point);
-  const xDistance = MarchParams.sampleSdf.value(offX);
-  const yDistance = MarchParams.sampleSdf.value(offY);
-  const zDistance = MarchParams.sampleSdf.value(offZ);
-
-  return normalize(
-    mul(
-      1 / epsilon,
-      vec3f(
-        xDistance - centerDistance,
-        yDistance - centerDistance,
-        zDistance - centerDistance,
-      ),
-    ),
+)((point) => {
+  const centerDistance = MarchParams.sampleSdf.$(point);
+  const distance = vec3f(
+    MarchParams.sampleSdf.$(point.add(epsilon.$.xyy)),
+    MarchParams.sampleSdf.$(point.add(epsilon.$.yxy)),
+    MarchParams.sampleSdf.$(point.add(epsilon.$.yyx)),
   );
+
+  return normalize(distance.sub(centerDistance));
 });

apps/phoure-www/src/lib-ray-marching/marchSdf.ts

@@ -3,55 +3,48 @@ import * as d from 'typegpu/data';
 import * as std from 'typegpu/std';
 import { ShapeContext } from './types.ts';
 
-type SampleSdf = TgpuFn<[d.Vec3f], d.F32>;
-
-const defaultGetSurfaceThreshold = tgpu['~unstable'].fn([ShapeContext], d.f32)(
-  (_ctx) => 0.001,
-);
+type SampleSdf = TgpuFn<(pos: d.Vec3f) => d.F32>;
 
 export const MarchParams = {
-  maxSteps: tgpu['~unstable'].slot(500),
+  maxSteps: tgpu.slot(500),
 
   /**
    * The distance from the camera at which the sky should be drawn instead of the world.
    */
-  farPlane: tgpu['~unstable'].slot(100),
+  farPlane: tgpu.slot(100),
 
-  getSurfaceThreshold: tgpu['~unstable'].slot(defaultGetSurfaceThreshold),
-  sampleSdf: tgpu['~unstable'].slot<SampleSdf>(),
+  surfaceThreshold: tgpu.slot<number>(0.001),
+  sampleSdf: tgpu.slot<SampleSdf>(),
 };
 
 export const MarchResult = d.struct({
   steps: d.u32,
   position: d.vec3f,
 });
 
-export const march = tgpu['~unstable'].fn([
+export const march = tgpu.fn([
   d.ptrFn(ShapeContext),
   d.u32,
   d.ptrFn(MarchResult),
 ])((ctx, limit, out) => {
   let pos = d.vec3f(ctx.rayPos);
   let prev_dist = d.f32(-1);
-  let min_dist = d.f32(MarchParams.farPlane.value);
+  let min_dist = d.f32(MarchParams.farPlane.$);
 
   let step = d.u32(0);
   let progress = d.f32(0);
 
   for (; step <= limit; step++) {
     pos = std.add(ctx.rayPos, std.mul(ctx.rayDir, progress));
-    min_dist = MarchParams.sampleSdf.value(pos);
+    min_dist = MarchParams.sampleSdf.$(pos);
 
     // Inside volume?
-    if (min_dist <= 0.) {
+    if (min_dist <= 0) {
       // No need to check more objects.
       break;
     }
 
-    if (
-      min_dist < MarchParams.getSurfaceThreshold.value(ctx) &&
-      min_dist < prev_dist
-    ) {
+    if (min_dist < MarchParams.surfaceThreshold.$ && min_dist < prev_dist) {
       // No need to check more objects.
       break;
     }
@@ -60,7 +53,7 @@ export const march = tgpu['~unstable'].fn([
     progress += min_dist;
     ctx.rayDistance += min_dist;
 
-    if (progress > MarchParams.farPlane.value) {
+    if (progress > MarchParams.farPlane.$) {
       // Stop checking.
       break;
     }
@@ -71,12 +64,9 @@ export const march = tgpu['~unstable'].fn([
   out.position = pos;
 
   // Not near surface or distance rising?
-  if (
-    min_dist > MarchParams.getSurfaceThreshold.value(ctx) * 2. ||
-    min_dist > prev_dist
-  ) {
+  if (min_dist > MarchParams.surfaceThreshold.$ * 2 || min_dist > prev_dist) {
     // Sky
-    out.steps = MarchParams.maxSteps.value + 1;
+    out.steps = MarchParams.maxSteps.$ + 1;
     return;
   }
 

apps/phoure-www/src/lib/GameEngine/sdfRenderer/sdfRenderer.ts

@@ -86,10 +86,7 @@ const renderSubPixel = tgpu.fn(
     return std.min(skyColor(init_shape_ctx.rayDir), d.vec3f(1));
   }
 
-  const init_normal = estimateNormal(
-    init_march_result.position,
-    init_shape_ctx,
-  );
+  const init_normal = estimateNormal(init_march_result.position);
   const init_material = Material();
 
   worldMat(init_march_result.position, init_shape_ctx, init_material);
@@ -135,7 +132,7 @@ const renderSubPixel = tgpu.fn(
         break;
       }
 
-      normal = estimateNormal(shape_ctx.rayPos, shape_ctx);
+      normal = estimateNormal(shape_ctx.rayPos);
 
       worldMat(shape_ctx.rayPos, shape_ctx, material);
 
@@ -236,7 +233,7 @@ const auxComputeFn = tgpu['~unstable'].computeFn({
   if (marchResult.steps >= MarchParams.maxSteps.$) {
     worldNormal = std.neg(shapeCtx.rayDir);
   } else {
-    worldNormal = estimateNormal(marchResult.position, shapeCtx);
+    worldNormal = estimateNormal(marchResult.position);
   }
 
   const material = Material();

apps/phoure-www/src/lib/GameEngine/sdfRenderer/worldSdf.ts

@@ -77,7 +77,6 @@ export const worldMat = tgpu.fn([d.vec3f, ShapeContext, d.ptrFn(Material)])((
   ctx,
   out,
 ) => {
-  const sd = MarchParams.getSurfaceThreshold.value(ctx);
   const d_left_blob = objLeftBlob(pos);
   const d_center_blob = objCenterBlob(pos);
   const d_right_blob = objRightBlob(pos);
@@ -87,21 +86,20 @@ export const worldMat = tgpu.fn([d.vec3f, ShapeContext, d.ptrFn(Material)])((
   out.emissive = false;
   out.roughness = 1;
 
-  if (d_left_blob <= sd) {
+  if (d_left_blob <= MarchParams.surfaceThreshold.$) {
     // left blob
     out.albedo = d.vec3f(1, 0.5, 0.2);
     out.roughness = 0.95;
-  } else if (d_center_blob <= sd) {
+  } else if (d_center_blob <= MarchParams.surfaceThreshold.$) {
     // test light
     out.albedo = mul(20, d.vec3f(1, 1, 0.5));
     out.emissive = true;
-  } else if (d_right_blob <= sd) {
+  } else if (d_right_blob <= MarchParams.surfaceThreshold.$) {
     out.albedo = mul(0.9, d.vec3f(0.5, 0.5, 0.6));
     out.roughness = 0.1;
-  } else if (d_floor_blob <= sd) {
+  } else if (d_floor_blob <= MarchParams.surfaceThreshold.$) {
     matFloor(pos, out);
   } else {
-    // out.albedo = vec3f(0.5, 0.5, 0.2);
     out.albedo = skyColor(ctx.rayDir);
   }
 });

pnpm-lock.yaml

@@ -4,5 +4,5 @@ packages:
 
 # Define a catalog of version ranges.
 catalog:
-  typegpu: https://pkg.pr.new/software-mansion/TypeGPU/typegpu@c81e38bfa323506e5933b26e7e5965450c68217b
+  typegpu: https://pkg.pr.new/software-mansion/TypeGPU/typegpu@e598b31a3d03476d47bd66002e80dfacb54d8203
   unplugin-typegpu: ^0.2.2