Fold: add hemi octahedron to plane operator

1 files changed, 58 insertions, 0 deletions

diff --git a/vertex_deformation.slang b/vertex_deformation.slang
index 50c252e..6e916ed 100644
--- a/vertex_deformation.slang
+++ b/vertex_deformation.slang
@@ -625,4 +625,62 @@ public void plane_to_hemi_octahedron_undeform_normal(float3 xyz, inout float3 no
   R3R3_UNDEFORM_NORMAL_AND_TANGENT(normal, tangent);
 }
 
+// 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
+[Differentiable]
+public float3 hemi_octahedron_to_plane(float3 xyz,
+    no_diff float3 p, no_diff float3 r_cart, no_diff float3 s_cart,
+    no_diff float t) {
+  // Convert from cartesian to (r, s, r x s) space.
+  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);
+
+  // Translate origin to `p` then change into (r, s, r x s) basis.
+  xyz = mul(to_rsrxs, xyz - p);
+
+  float3 xyz0 = xyz;
+
+  // Undo the -45° unrotation from forward pass (rotate by +45°)
+  float x_rot = (xyz.x + xyz.z) * RCP_SQRT_2;
+  float z_rot = (xyz.z - xyz.x) * RCP_SQRT_2;
+
+  // Octahedral encode: project normalized direction to plane
+  float3 N = normalize(float3(x_rot, xyz.y, z_rot));
+  N.y = dmax(N.y, 1e-4f);
+  float L1 = dabs(N.x) + dabs(N.y) + dabs(N.z);
+  float x_oct = N.x / L1;
+  float z_oct = N.z / L1;
+
+  // Undo the initial 45° rotation (x_rot = (x+z)*0.5, z_rot = (z-x)*0.5)
+  // Inverse: x = x_rot - z_rot, z = x_rot + z_rot
+  float x_plane = x_oct - z_oct;
+  float z_plane = x_oct + z_oct;
+
+  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));
+
+  // Map back to cartesian basis
+  xyz = mul(to_cart, result) + p;
+
+  return xyz;
+}
+
+public void hemi_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, hemi_octahedron_to_plane, p, r, s, t);
+}
+
+public void hemi_octahedron_to_plane_undeform_normal(float3 xyz, inout float3 normal,
+    inout float3 tangent, float3 p, float3 r, float3 s, float t) {
+  R3R3_DECLARE_BASIS_VECTORS(xyz);
+  R3R3_AUTODIFF_BASIS_VECTORS(hemi_octahedron_to_plane, p, r, s, t);
+  R3R3_UNDEFORM_NORMAL_AND_TANGENT(normal, tangent);
+}
+
 #endif  // __CUSTOM31_INC