Add constant, linear and TCB spline formula operations

Co-Authored-By: Claude Haiku 5.5 <noreply@anthropic.com>

1 files changed, 40 insertions, 2 deletions

diff --git a/generat3d/formula.py b/generat3d/formula.py
index 67d25e0..2165e16 100644
--- a/generat3d/formula.py
+++ b/generat3d/formula.py
@@ -15,6 +15,7 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
+import bisect
 import math
 import operator
 
@@ -23,7 +24,7 @@ from . import Error
 STAGES = ("sum", "mult")
 BINARY = {"add": operator.add, "sub": operator.sub,
           "mult": operator.mul, "div": operator.truediv}
-SPLINES = ()  # filled in by Task 4
+SPLINES = ("spline_constant", "spline_linear", "spline_tcb")
 
 
 class Prop:
@@ -122,4 +123,41 @@ class Graph:
 
 
 def _spline_op(name, stack):
-    raise NotImplementedError  # Task 4
+    n = int(stack.pop())
+    if n < 1 or n > len(stack) - 1:
+        raise IndexError
+    knots = stack[-n:]
+    del stack[-n:]
+    return spline(name, stack.pop(), knots)
+
+
+def spline(kind, x, knots):
+    """Evaluate a DSON spline at x. Knots are [x, y] or [x, y, t, c, b], sorted by x."""
+    xs = [k[0] for k in knots]
+    ys = [k[1] for k in knots]
+    if x <= xs[0]:
+        return ys[0]
+    if x >= xs[-1]:
+        return ys[-1]
+    i = bisect.bisect_right(xs, x) - 1
+    if kind == "spline_constant":
+        return ys[i]
+    s = (x - xs[i]) / (xs[i + 1] - xs[i])
+    if kind == "spline_linear":
+        return ys[i] + s * (ys[i + 1] - ys[i])
+    # ponytail: Kochanek-Bartels with uniform segment parameterization;
+    # if the corpus shows Daz scales tangents by knot spacing, add that here.
+    m0 = _tcb_tangent(knots, ys, i, outgoing=True)
+    m1 = _tcb_tangent(knots, ys, i + 1, outgoing=False)
+    s2, s3 = s * s, s * s * s
+    return ((2 * s3 - 3 * s2 + 1) * ys[i] + (s3 - 2 * s2 + s) * m0
+            + (-2 * s3 + 3 * s2) * ys[i + 1] + (s3 - s2) * m1)
+
+
+def _tcb_tangent(knots, ys, k, outgoing):
+    t, c, b = (list(knots[k][2:5]) + [0, 0, 0])[:3]
+    d_in = ys[k] - ys[max(k - 1, 0)]
+    d_out = ys[min(k + 1, len(ys) - 1)] - ys[k]
+    if outgoing:
+        return (1 - t) * ((1 + b) * (1 + c) * d_in + (1 - b) * (1 - c) * d_out) / 2
+    return (1 - t) * ((1 + b) * (1 - c) * d_in + (1 - b) * (1 + c) * d_out) / 2