add "axis align" feature to fix distortion after unbending the tube

1 files changed, 52 insertions, 1 deletions

diff --git a/vertex_deformation.slang b/vertex_deformation.slang
index aea1b5b..0846991 100644
--- a/vertex_deformation.slang
+++ b/vertex_deformation.slang
@@ -106,7 +106,8 @@ public float3 plane_to_tube(float3 xyz,
   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  = inverse(to_rsrxs, determinant(to_rsrxs));
+  // Ivnerse of orthonormal matrix is just the transpose.
+  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);
@@ -155,6 +156,56 @@ public void plane_to_tube_undeform_normal(float3 xyz, inout float3 normal,
   R3R3_UNDEFORM_NORMAL_AND_TANGENT(normal, tangent);
 }
 
+[Differentiable]
+public float3 axis_align(float3 xyz, no_diff float3 po, no_diff float3 pp, no_diff float3 r, no_diff float t) {
+  // We assume that `s` is orthogonal to `r`, and that `r` is normalized.
+  float3 s = normalize(pp - po);
+  float3 sxr = cross(s, r);
+
+  float3x3 to_rsrxs = float3x3(r, s, sxr);
+  // Inverse of orthonormal matrix is just the transpose.
+  float3x3 to_cart  = transpose(to_rsrxs);
+
+  // Move key vectors into (r, s, sxr) basis centered at `po`, per the derivation.
+  float3 xyz_rsrxs = mul(to_rsrxs, xyz - po);
+
+  float vr = xyz_rsrxs[0];
+  float vs = xyz_rsrxs[1];
+  float2 v0 = float2(vr, vs);
+  float3 pp_rsrxs = mul(to_rsrxs, pp - po);
+  float qr = pp_rsrxs[1] * vr / vs;  // TODO epsilon
+  float2 q = float2(qr, pp_rsrxs[1]);
+
+  // Translate to `q`.
+  v0 -= q;
+
+  // Project onto `s`.
+  v0.x = 0;
+
+  // Translate back to `po`.
+  v0 += q;
+
+  xyz_rsrxs[0] = v0[0];
+  xyz_rsrxs[1] = v0[1];
+
+  // Move back into cartesian space.
+  xyz = mul(to_cart, xyz_rsrxs) + po;
+
+  return xyz;
+}
+
+public void axis_align_normal(inout float3 xyz, inout float3 normal,
+    inout float3 tangent, float3 po, float3 pp, float3 r, float t) {
+  R3R3_NORMALS(xyz, normal, tangent, axis_align, po, pp, r, t);
+}
+
+public void axis_align_undeform_normal(float3 xyz, inout float3 normal,
+    inout float3 tangent, float3 po, float3 pp, float3 r, float t) {
+  R3R3_DECLARE_BASIS_VECTORS(xyz);
+  R3R3_AUTODIFF_BASIS_VECTORS(axis_align, po, pp, r, t);
+  R3R3_UNDEFORM_NORMAL_AND_TANGENT(normal, tangent);
+}
+
 // Maps a tube with circular cross section on the xz plane to a quad on the xy
 // plane.
 [Differentiable]