Fold: add plane to octahedron code

1 files changed, 89 insertions, 9 deletions

diff --git a/vertex_deformation.slang b/vertex_deformation.slang
index 0883c84..de53394 100755
--- a/vertex_deformation.slang
+++ b/vertex_deformation.slang
@@ -483,10 +483,7 @@ public void fbm_normal(inout float3 xyz, inout float3 normal, inout float3 tange
   tangent = mul(jac, tangent);
 }
 
-// Maps a plane to a hemi-octahedron (half octahedron).
-// Uses octahedral parameterization consistent with pbrt's OctahedralVector.
-// Input: plane on [-1,1]² in the (r, rxs) plane
-// Output: unit hemisphere with pole at +s direction
+// Maps a plane on [-1, 1] on xz plane to a hemi-octahedron with radius 1.
 [Differentiable]
 public float3 plane_to_hemi_octahedron(float3 xyz,
     no_diff float3 p, no_diff float3 r_cart, no_diff float3 s_cart,
@@ -526,7 +523,7 @@ public float3 plane_to_hemi_octahedron(float3 xyz,
   oct_pos = float3(x_unrot, oct_pos.y, z_unrot);
 
   // Interpolate between original position and sphere position
-  float3 result = dlerp(xyz0, oct_pos, dmin(t, 1.0f));
+  float3 result = dlerp(xyz0, oct_pos, t);
 
   // Map back to cartesian basis
   xyz = mul(to_cart, result) + p;
@@ -539,9 +536,7 @@ public void plane_to_hemi_octahedron_normal(inout float3 xyz, inout float3 norma
   R3R3_NORMALS(xyz, normal, tangent, plane_to_hemi_octahedron, p, r, s, t);
 }
 
-// Maps a hemi-octahedron to a plane (inverse of plane_to_hemi_octahedron).
-// Input: unit hemisphere with pole at +s direction
-// Output: plane on [-1,1]² in the (r, rxs) plane
+// Maps a hemi-octahedron with raidus 1 to a quad on [-1, 1] on the (r, rxs) plane.
 [Differentiable]
 public float3 hemi_octahedron_to_plane(float3 xyz,
     no_diff float3 p, no_diff float3 r_cart, no_diff float3 s_cart,
@@ -577,7 +572,7 @@ public float3 hemi_octahedron_to_plane(float3 xyz,
   float3 plane_pos = float3(x_plane, 0.0f, z_plane);
 
   // Interpolate between original position and plane position
-  float3 result = dlerp(xyz0, plane_pos, dmin(t, 1.0f));
+  float3 result = dlerp(xyz0, plane_pos, t);
 
   // Map back to cartesian basis
   xyz = mul(to_cart, result) + p;
@@ -590,6 +585,73 @@ public void hemi_octahedron_to_plane_normal(inout float3 xyz, inout float3 norma
   R3R3_NORMALS(xyz, normal, tangent, hemi_octahedron_to_plane, p, r, s, t);
 }
 
+// Maps [-1, 1] on (r, rxs) plane to a unit sphere using octahedral mapping.
+[Differentiable]
+public float3 plane_to_octahedron(float3 xyz,
+    no_diff float3 p, no_diff float3 r_cart, no_diff float3 s_cart,
+    no_diff float t) {
+  r_cart = normalize(r_cart);
+  s_cart = normalize(s_cart);
+  float3 rxs_cart = cross(s_cart, r_cart);
+  float3x3 to_rsrxs = float3x3(r_cart, s_cart, rxs_cart);
+  float3x3 to_cart = transpose(to_rsrxs);
+
+  xyz = mul(to_rsrxs, xyz - p);
+  float3 xyz0 = xyz;
+
+  float l1_norm = dabs(xyz.x) + dabs(xyz.z);
+  if (l1_norm > 1) {
+    xyz.x = sign(xyz0.x) * (1 - dabs(xyz0.z));
+    xyz.z = sign(xyz0.z) * (1 - dabs(xyz0.x));
+  }
+  xyz.y = 1 - l1_norm;
+
+  xyz *= (1 + xyz0.y);
+
+  float3 result = dlerp(xyz0, xyz, t);
+
+  xyz = mul(to_cart, result) + p;
+  return xyz;
+}
+
+public void plane_to_octahedron_normal(inout float3 xyz, inout float3 normal,
+    inout float3 tangent, float3 p, float3 r, float3 s, float t) {
+  R3R3_NORMALS(xyz, normal, tangent, plane_to_octahedron, p, r, s, t);
+}
+
+// Maps a unit sphere to a plane using octahedral mapping.
+[Differentiable]
+public float3 octahedron_to_plane(float3 xyz,
+    no_diff float3 p, no_diff float3 r_cart, no_diff float3 s_cart,
+    no_diff float t) {
+  r_cart = normalize(r_cart);
+  s_cart = normalize(s_cart);
+  float3 rxs_cart = cross(s_cart, r_cart);
+  float3x3 to_rsrxs = float3x3(r_cart, s_cart, rxs_cart);
+  float3x3 to_cart = transpose(to_rsrxs);
+
+  xyz = mul(to_rsrxs, xyz - p);
+  float3 xyz0 = xyz;
+
+  float l1_norm = dabs(xyz.x) + dabs(xyz.y) + dabs(xyz.z);
+  xyz /= l1_norm;
+  float2 xz_tmp = xyz.xz;
+  if (xyz.y < 0) {
+    xyz.x = sign(xz_tmp[0]) * (1 - dabs(xz_tmp[1]));
+    xyz.z = sign(xz_tmp[1]) * (1 - dabs(xz_tmp[0]));
+  }
+  xyz.y = 0;
+
+  float3 result = dlerp(xyz0, xyz, t);
+  xyz = mul(to_cart, result) + p;
+  return xyz;
+}
+
+public void octahedron_to_plane_normal(inout float3 xyz, inout float3 normal,
+    inout float3 tangent, float3 p, float3 r, float3 s, float t) {
+  R3R3_NORMALS(xyz, normal, tangent, octahedron_to_plane, p, r, s, t);
+}
+
 [Differentiable]
 public float3 scale(float3 xyz,
     no_diff float3 k, no_diff float t) {
@@ -612,5 +674,23 @@ public void translate_normal(inout float3 xyz, inout float3 normal,
   R3R3_NORMALS(xyz, normal, tangent, translate, offset, t);
 }
 
+[Differentiable]
+public float3 rotate(float3 xyz,
+    no_diff float3 p, no_diff float3 axis, no_diff float angle, no_diff float t) {
+  float theta = angle * t;
+  float3 a = normalize(axis);
+  float c = cos(theta);
+  float s = sin(theta);
+  float3 v = xyz - p;
+  // Rodrigues' rotation formula
+  float3 rotated = v * c + cross(a, v) * s + a * dot(a, v) * (1.0f - c);
+  return rotated + p;
+}
+
+public void rotate_normal(inout float3 xyz, inout float3 normal,
+    inout float3 tangent, float3 p, float3 axis, float angle, float t) {
+  R3R3_NORMALS(xyz, normal, tangent, rotate, p, axis, angle, t);
+}
+
 #endif  // __CUSTOM31_INC