Add constant, linear and TCB spline formula operations

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

2 files changed, 78 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
diff --git a/tests/test_formula.py b/tests/test_formula.py
index 72f41ea..7440681 100644
--- a/tests/test_formula.py
+++ b/tests/test_formula.py
@@ -130,3 +130,41 @@ class GraphTest(unittest.TestCase):
         g.add_formula("c", [push("a"), push("a"), {"op": "add"}])
         g.add_formula("d", [push("b"), push("c"), {"op": "add"}])
         self.assertEqual(g.evaluate()["d"], 3)
+
+
+from generat3d.formula import spline
+
+
+def spline_ops(kind, knots):
+    return ([{"op": "push", "key": "a"}] + [val(k) for k in knots]
+            + [val(len(knots)), {"op": kind}])
+
+
+class SplineTest(unittest.TestCase):
+    LINE = [[0, 0], [1, 1], [2, 2], [3, 3]]
+
+    def test_outside_range_holds_ends(self):
+        for kind in ("spline_constant", "spline_linear", "spline_tcb"):
+            self.assertEqual(spline(kind, -1, self.LINE), 0)
+            self.assertEqual(spline(kind, 9, self.LINE), 3)
+
+    def test_constant(self):
+        self.assertEqual(spline("spline_constant", 1.7, [[0, 5], [1, 6], [2, 7]]), 6)
+
+    def test_linear(self):
+        self.assertAlmostEqual(spline("spline_linear", 0.25, [[0, 0], [1, 4]]), 1.0)
+
+    def test_tcb_hits_knots(self):
+        knots = [[0, 0, 0, 0, 0], [1, 3, 0, 0, 0], [2, 1, 0, 0, 0]]
+        self.assertEqual(spline("spline_tcb", 1, knots), 3)
+
+    def test_tcb_zero_params_is_catmull_rom(self):
+        knots = [k + [0, 0, 0] for k in self.LINE]
+        self.assertAlmostEqual(spline("spline_tcb", 1.5, knots), 1.5)
+
+    def test_spline_in_graph(self):
+        g = Graph()
+        g.add_prop("a", 0.5)
+        g.add_prop("b")
+        g.add_formula("b", spline_ops("spline_linear", [[0, 0], [1, 10]]))
+        self.assertAlmostEqual(g.evaluate()["b"], 5.0)