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

5 files changed, 94 insertions, 1 deletions

diff --git a/3ner.shader b/3ner.shader
index 766091c..eadde89 100644
--- a/3ner.shader
+++ b/3ner.shader
@@ -193,6 +193,16 @@ Shader "yum_food/3ner"
             [HideInInspector] m_end_Vertex_Deformation_Tube_To_Plane_0("Tube to Plane 0", Float) = 0
             //endex
 
+            //ifex _Vertex_Deformation_Axis_Align_Enabled==0
+            [HideInInspector] m_start_Vertex_Deformation_Axis_Align("Axis Align", Float) = 0
+              [ThryToggle(_VERTEX_DEFORMATION_AXIS_ALIGN)] _Vertex_Deformation_Axis_Align_Enabled("Enable", Float) = 0
+              _Vertex_Deformation_Axis_Align_po("po", Vector) = (0, 0, 0)
+              _Vertex_Deformation_Axis_Align_pp("pp", Vector) = (0, 0, 1)
+              _Vertex_Deformation_Axis_Align_r("r", Vector) = (0, 1, 0)
+              _Vertex_Deformation_Axis_Align_t("t", Range(0,1)) = 0
+            [HideInInspector] m_end_Vertex_Deformation_Axis_Align("Axis Align", Float) = 0
+            //endex
+
             //ifex _Vertex_Deformation_Tube_To_Plane_1_Enabled==0
             [HideInInspector] m_start_Vertex_Deformation_Tube_To_Plane_1("Tube to Plane 1", Float) = 0
               [ThryToggle(_VERTEX_DEFORMATION_TUBE_TO_PLANE_1)] _Vertex_Deformation_Tube_To_Plane_1_Enabled("Enable", Float) = 0
diff --git a/features.cginc b/features.cginc
index 71fb9b6..6dbff87 100644
--- a/features.cginc
+++ b/features.cginc
@@ -40,6 +40,7 @@
 //ifex _Vertex_Deformation_Tubes_Enabled==0
 #pragma shader_feature_local _VERTEX_DEFORMATION_TUBES
 #pragma shader_feature_local _VERTEX_DEFORMATION_TUBE_TO_PLANE_0
+#pragma shader_feature_local _VERTEX_DEFORMATION_AXIS_ALIGN
 #pragma shader_feature_local _VERTEX_DEFORMATION_TUBE_TO_PLANE_1
 #pragma shader_feature_local _VERTEX_DEFORMATION_PLANE_TO_TUBE_0
 #pragma shader_feature_local _VERTEX_DEFORMATION_PLANE_TO_TUBE_1
diff --git a/globals.cginc b/globals.cginc
index b5d6b21..21d9889 100644
--- a/globals.cginc
+++ b/globals.cginc
@@ -82,6 +82,13 @@ float3 _Vertex_Deformation_Tube_To_Plane_0_s;
 float _Vertex_Deformation_Tube_To_Plane_0_t;
 #endif  // _VERTEX_DEFORMATION_TUBE_TO_PLANE_0
 
+#if defined(_VERTEX_DEFORMATION_AXIS_ALIGN)
+float3 _Vertex_Deformation_Axis_Align_po;
+float3 _Vertex_Deformation_Axis_Align_pp;
+float3 _Vertex_Deformation_Axis_Align_r;
+float _Vertex_Deformation_Axis_Align_t;
+#endif  // _VERTEX_DEFORMATION_AXIS_ALIGN
+
 #if defined(_VERTEX_DEFORMATION_TUBE_TO_PLANE_1)
 float3 _Vertex_Deformation_Tube_To_Plane_1_p;
 float3 _Vertex_Deformation_Tube_To_Plane_1_r;
diff --git a/vertex.cginc b/vertex.cginc
index 3d87f35..67565d9 100644
--- a/vertex.cginc
+++ b/vertex.cginc
@@ -22,6 +22,22 @@
 #define VERTEX_DEFORM_TUBE_TO_PLANE_0_NORM
 #endif  // VERTEX_DEFORMATION_TUBE_TO_PLANE_0
 
+#if defined(_VERTEX_DEFORMATION_AXIS_ALIGN)
+#define VERTEX_DEFORM_AXIS_ALIGN_PREAMBLE                    \
+  float3 po = _Vertex_Deformation_Axis_Align_po;               \
+  float3 pp = _Vertex_Deformation_Axis_Align_pp;               \
+  float3 r = _Vertex_Deformation_Axis_Align_r;               \
+  float t = _Vertex_Deformation_Axis_Align_t
+#define VERTEX_DEFORM_AXIS_ALIGN_POS                         \
+  objPos = axis_align(objPos, po, pp, r, t)
+#define VERTEX_DEFORM_AXIS_ALIGN_NORM                        \
+  axis_align_normal(objPos, objNorm, objTan, po, pp, r, t)
+#else
+#define VERTEX_DEFORM_AXIS_ALIGN_PREAMBLE
+#define VERTEX_DEFORM_AXIS_ALIGN_POS
+#define VERTEX_DEFORM_AXIS_ALIGN_NORM
+#endif  // VERTEX_DEFORMATION_AXIS_ALIGN
+
 #if defined(_VERTEX_DEFORMATION_TUBE_TO_PLANE_1)
 #define VERTEX_DEFORM_TUBE_TO_PLANE_1_PREAMBLE                    \
   float3 p = _Vertex_Deformation_Tube_To_Plane_1_p;               \
@@ -143,6 +159,10 @@ void deform(inout float3 objPos) {
     VERTEX_DEFORM_TUBE_TO_PLANE_0_POS;
   }
   {
+    VERTEX_DEFORM_AXIS_ALIGN_PREAMBLE;
+    VERTEX_DEFORM_AXIS_ALIGN_POS;
+  }
+  {
     VERTEX_DEFORM_TUBE_TO_PLANE_1_PREAMBLE;
     VERTEX_DEFORM_TUBE_TO_PLANE_1_POS;
   }
@@ -179,6 +199,10 @@ void deform_normal(inout float3 objPos, inout float3 objNorm, inout float3 objTa
     VERTEX_DEFORM_TUBE_TO_PLANE_0_NORM;
   }
   {
+    VERTEX_DEFORM_AXIS_ALIGN_PREAMBLE;
+    VERTEX_DEFORM_AXIS_ALIGN_NORM;
+  }
+  {
     VERTEX_DEFORM_TUBE_TO_PLANE_1_PREAMBLE;
     VERTEX_DEFORM_TUBE_TO_PLANE_1_NORM;
   }
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]