solving splines bugs

pina/model/spline.py

@@ -6,26 +6,26 @@
 
 class Spline(torch.nn.Module):
 
-    def __init__(self, order, knots=None, control_points=None) -> None:
+    def __init__(self, order=4, knots=None, control_points=None) -> None:
         """
         Spline model.
 
         :param int order: the order of the spline.
 
         """
         super().__init__()
-
         check_consistency(order, int)
         if order < 0:
             raise ValueError("Spline order cannot be negative.")
         if knots is None and control_points is None:
             raise ValueError("Knots and control points cannot be both None.")
 
         self.order = order
+        self.k = order-1
 
         if knots is not None:
             self.knots = knots
-            n = len(knots) - order - 1
+            n = len(knots) - order
             self.control_points = torch.nn.Parameter(
                 torch.zeros(n), requires_grad=True)
 
@@ -36,13 +36,12 @@ def __init__(self, order, knots=None, control_points=None) -> None:
                 'type': 'auto',
                 'min': 0,
                 'max': 1,
-                'n': n + order + 1}
+                'n': n}
 
         else:
             self.knots = knots
             self.control_points = control_points
 
-        print(self.knots)
 
         if self.knots.ndim != 1:
             raise ValueError("Knot vector must be one-dimensional.")
@@ -93,7 +92,7 @@ def __init__(self, order, knots=None, control_points=None) -> None:
         # plt.show()
 
     @staticmethod
-    def B(x, k, i, t, order):
+    def B(x, k, i, t):
         '''
         x: points to be evaluated
         k: spline degree
@@ -102,21 +101,21 @@ def B(x, k, i, t, order):
         '''
         if k == 0:
             a = torch.where(torch.logical_and(t[i] <= x, x < t[i+1]), 1.0, 0.0)
-            if i == len(t) - order - 1:
-                a = torch.where(x == t[-1], 1.0, a)
+            if i == len(t) - k - 1:
+                 a = torch.where(x == t[-1], 1.0, a)
             a.requires_grad_(True)
             return a
 
 
         if t[i+k] == t[i]:
             c1 = torch.tensor([0.0], requires_grad=True)
         else:
-            c1 = (x - t[i])/(t[i+k] - t[i]) * Spline.B(x, k-1, i, t, order=order)
+            c1 = (x - t[i])/(t[i+k] - t[i]) * Spline.B(x, k-1, i, t)
 
         if t[i+k+1] == t[i+1]:
             c2 = torch.tensor([0.0], requires_grad=True)
         else:
-            c2 = (t[i+k+1] - x)/(t[i+k+1] - t[i+1]) * Spline.B(x, k-1, i+1, t, order=order)
+            c2 = (t[i+k+1] - x)/(t[i+k+1] - t[i+1]) * Spline.B(x, k-1, i+1, t)
 
         return c1 + c2
 
@@ -134,7 +133,6 @@ def control_points(self, value):
             dim = value.get('dim', 1)
             value = torch.zeros(n, dim)
 
-        print(value)
         if not isinstance(value, torch.Tensor):
             raise ValueError('Invalid value for control_points')
 
@@ -154,14 +152,13 @@ def knots(self, value):
             n = value.get('n', 10)
 
             if type_ == 'uniform':
-                value = torch.linspace(min_, max_, n + k + 1)
+                value = torch.linspace(min_, max_, n + self.k + 1)
             elif type_ == 'auto':
-                k = self.order - 1
                 value = torch.concatenate(
                     (
-                        torch.ones(k)*min_,
-                        torch.linspace(min_, max_, n - k +1),
-                        torch.ones(k)*max_,
+                        torch.ones(self.k)*min_,
+                        torch.linspace(min_, max_, n - self.k +1),
+                        torch.ones(self.k)*max_,
                         # [self.max] * (k-1)
                     )
                 )
@@ -173,7 +170,7 @@ def knots(self, value):
 
     def forward(self, x_):
         t = self.knots
-        k = self.order
+        k = self.k
         c = self.control_points
 
         # return LabelTensor((x_**2).reshape(-1, 1), ['v'])
@@ -186,7 +183,7 @@ def forward(self, x_):
         #     print(self.B(x_, k, i, t), c[i])
         #     print(self.B(x, k, i, t), c[i])
         tmp_result = torch.concatenate([
-            (c[i] * Spline.B(x_, k, i, t, self.order)).reshape(
+            (c[i] * Spline.B(x_, k, i, t)).reshape(
                 1, x_.shape[0], -1) 
             for i in range(len(c))], axis=0
         )

tests/test_model/test_spline.py

@@ -19,12 +19,11 @@ def scipy_check(model, x, y):
     spline = BSpline(
         t=model.knots.detach(),
         c=model.control_points.detach(),
-        k=model.order
+        k=model.order-1
     )
-    y_numpy = spline(x)
+    y_scipy = spline(x)
     y = y.detach().numpy()
-    print(y)
-    np.testing.assert_allclose(y, y_numpy, atol=1e-5)
+    np.testing.assert_allclose(y, y_scipy, atol=1e-5)
 
 def test_constructor():
     Spline()
@@ -37,9 +36,10 @@ def test_constructor_wrong():
 @pytest.mark.parametrize("args", valid_args)
 def test_forward(args):
     xi = torch.linspace(0, 1, 100)
-    model = Spline(**args)
+    model = Spline(*args)
     yi = model(xi).squeeze()
     scipy_check(model, xi, yi)
+    return 
 
 
 def test_backward():