Fold: more animation bugfixes

1 files changed, 17 insertions, 12 deletions

diff --git a/vertex_deformation.slang b/vertex_deformation.slang
index de53394..b02d154 100755
--- a/vertex_deformation.slang
+++ b/vertex_deformation.slang
@@ -330,10 +330,10 @@ public void sine_wave_normal(inout float3 xyz, inout float3 normal, inout float3
 
 [Differentiable]
 public float3 norm_conversion(float3 xyz, no_diff float input_k, no_diff float output_k, no_diff float t) {
-  float3 xyz_abs = abs(xyz)+1e-4f;
+  float3 xyz_abs = dabs(xyz)+1e-4f;
   float lin = pow(pow(xyz_abs.x, input_k) + pow(xyz_abs.y, input_k) + pow(xyz_abs.z, input_k), 1.0f / input_k);
   float lout = pow(pow(xyz_abs.x, output_k) + pow(xyz_abs.y, output_k) + pow(xyz_abs.z, output_k), 1.0f / output_k);
-  float scale = lerp(1.0f, lout / lin, t);
+  float scale = dlerp(1.0f, lout / lin, t);
   return xyz * scale;
 }
 
@@ -512,10 +512,9 @@ public float3 plane_to_hemi_octahedron(float3 xyz,
   // Use differentiable abs and max for smooth autodiff
   float y = dmax(0.0f, 1.0f - dabs(x_rot) - dabs(z_rot));
 
-  // Normalize to unit sphere (differentiable safe normalize)
   float3 oct_pos = float3(x_rot, y, z_rot);
   float len = dot(oct_pos, oct_pos);
-  oct_pos = (oct_pos / sqrt(len)) * (1.0f + xyz.y);
+  oct_pos = oct_pos * (1.0f + xyz.y);
 
   // Rotate back by -45° around y to undo input rotation
   float x_unrot = (oct_pos.x - oct_pos.z) * RCP_SQRT_2;
@@ -557,19 +556,22 @@ public float3 hemi_octahedron_to_plane(float3 xyz,
   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;
+  // The forward pass scales by (1 + s_original), and the unscaled hemi-octahedron
+  // has L1 norm = 1, so L1 of the scaled point recovers s_original.
+  float L1_scaled = dabs(x_rot) + dabs(xyz.y) + dabs(z_rot);
+  float s_original = L1_scaled - 1.0f;
+
+  // Remove (1 + s) scale to get the unit hemi-octahedron point, then decode.
+  float L1_inv = 1.0f / L1_scaled;
+  float x_oct = x_rot * L1_inv;
+  float z_oct = z_rot * L1_inv;
 
   // 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);
+  float3 plane_pos = float3(x_plane, s_original, z_plane);
 
   // Interpolate between original position and plane position
   float3 result = dlerp(xyz0, plane_pos, t);
@@ -633,14 +635,17 @@ public float3 octahedron_to_plane(float3 xyz,
   xyz = mul(to_rsrxs, xyz - p);
   float3 xyz0 = xyz;
 
+  // The forward pass scales by (1 + s_original), and the unscaled octahedron
+  // has L1 norm = 1, so L1 of the scaled point recovers s_original.
   float l1_norm = dabs(xyz.x) + dabs(xyz.y) + dabs(xyz.z);
+  float s_original = l1_norm - 1.0f;
   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;
+  xyz.y = s_original;
 
   float3 result = dlerp(xyz0, xyz, t);
   xyz = mul(to_cart, result) + p;