add b-spline surface

tests/test_model/test_spline.py

@@ -1,6 +1,5 @@
 import torch
 import pytest
-import numpy as np
 from scipy.interpolate import BSpline
 from pina.model import Spline
 from pina import LabelTensor
@@ -12,7 +11,10 @@ n_ctrl_pts = torch.randint(order, order + 5, (1,)).item()
 n_knots = order + n_ctrl_pts
 
 # Input tensor
-pts = LabelTensor(torch.linspace(0, 1, 100).reshape(-1, 1), ["x"])
+points = [
+    LabelTensor(torch.rand(100, 1), ["x"]),
+    LabelTensor(torch.rand(2, 100, 1), ["x"]),
+]
 
 
 # Function to compare with scipy implementation
@@ -26,15 +28,15 @@ def check_scipy_spline(model, x, output_):
     )
 
     # Compare outputs
-    np.testing.assert_allclose(
-        output_.squeeze().detach().numpy(),
-        scipy_spline(x).flatten(),
+    torch.allclose(
+        output_,
+        torch.tensor(scipy_spline(x), dtype=output_.dtype),
         atol=1e-5,
         rtol=1e-5,
     )
 
 
-# Define all possible combinations of valid arguments for the Spline class
+# Define all possible combinations of valid arguments for Spline class
 valid_args = [
     {
         "order": order,
@@ -144,14 +146,15 @@ def test_constructor(args):
 
 
 @pytest.mark.parametrize("args", valid_args)
-def test_forward(args):
+@pytest.mark.parametrize("pts", points)
+def test_forward(args, pts):
 
     # Define the model
     model = Spline(**args)
 
     # Evaluate the model
     output_ = model(pts)
-    assert output_.shape == (pts.shape[0], 1)
+    assert output_.shape == pts.shape
 
     # Compare with scipy implementation only for interpolant knots (mode: auto)
     if isinstance(args["knots"], dict) and args["knots"]["mode"] == "auto":
@@ -159,7 +162,8 @@ def test_forward(args):
 
 
 @pytest.mark.parametrize("args", valid_args)
-def test_backward(args):
+@pytest.mark.parametrize("pts", points)
+def test_backward(args, pts):
 
     # Define the model
     model = Spline(**args)

tests/test_model/test_spline_surface.py

@@ -0,0 +1,180 @@
+import torch
+import random
+import pytest
+from pina.model import SplineSurface
+from pina import LabelTensor
+
+
+# Utility quantities for testing
+orders = [random.randint(1, 8) for _ in range(2)]
+n_ctrl_pts = random.randint(max(orders), max(orders) + 5)
+n_knots = [orders[i] + n_ctrl_pts for i in range(2)]
+
+# Input tensor
+points = [
+    LabelTensor(torch.rand(100, 2), ["x", "y"]),
+    LabelTensor(torch.rand(2, 100, 2), ["x", "y"]),
+]
+
+
+@pytest.mark.parametrize(
+    "knots_u",
+    [
+        torch.rand(n_knots[0]),
+        {"n": n_knots[0], "min": 0, "max": 1, "mode": "auto"},
+        {"n": n_knots[0], "min": 0, "max": 1, "mode": "uniform"},
+        None,
+    ],
+)
+@pytest.mark.parametrize(
+    "knots_v",
+    [
+        torch.rand(n_knots[1]),
+        {"n": n_knots[1], "min": 0, "max": 1, "mode": "auto"},
+        {"n": n_knots[1], "min": 0, "max": 1, "mode": "uniform"},
+        None,
+    ],
+)
+@pytest.mark.parametrize(
+    "control_points", [torch.rand(n_ctrl_pts, n_ctrl_pts), None]
+)
+def test_constructor(knots_u, knots_v, control_points):
+
+    # Skip if knots_u, knots_v, and control_points are all None
+    if (knots_u is None or knots_v is None) and control_points is None:
+        return
+
+    SplineSurface(
+        orders=orders,
+        knots_u=knots_u,
+        knots_v=knots_v,
+        control_points=control_points,
+    )
+
+    # Should fail if orders is not list of two elements
+    with pytest.raises(ValueError):
+        SplineSurface(
+            orders=[orders[0]],
+            knots_u=knots_u,
+            knots_v=knots_v,
+            control_points=control_points,
+        )
+
+    # Should fail if both knots and control_points are None
+    with pytest.raises(ValueError):
+        SplineSurface(
+            orders=orders,
+            knots_u=None,
+            knots_v=None,
+            control_points=None,
+        )
+
+    # Should fail if control_points is not a torch.Tensor when provided
+    with pytest.raises(ValueError):
+        SplineSurface(
+            orders=orders,
+            knots_u=knots_u,
+            knots_v=knots_v,
+            control_points=[[0.0] * n_ctrl_pts] * n_ctrl_pts,
+        )
+
+    # Should fail if control_points is not of the correct shape when provided
+    # It assumes that at least one among knots_u and knots_v is not None
+    if knots_u is not None or knots_v is not None:
+        with pytest.raises(ValueError):
+            SplineSurface(
+                orders=orders,
+                knots_u=knots_u,
+                knots_v=knots_v,
+                control_points=torch.rand(n_ctrl_pts + 1, n_ctrl_pts + 1),
+            )
+
+    # Should fail if there are not enough knots_u to define the control points
+    with pytest.raises(ValueError):
+        SplineSurface(
+            orders=orders,
+            knots_u=torch.linspace(0, 1, orders[0]),
+            knots_v=knots_v,
+            control_points=None,
+        )
+
+    # Should fail if there are not enough knots_v to define the control points
+    with pytest.raises(ValueError):
+        SplineSurface(
+            orders=orders,
+            knots_u=knots_u,
+            knots_v=torch.linspace(0, 1, orders[1]),
+            control_points=None,
+        )
+
+
+@pytest.mark.parametrize(
+    "knots_u",
+    [
+        torch.rand(n_knots[0]),
+        {"n": n_knots[0], "min": 0, "max": 1, "mode": "auto"},
+        {"n": n_knots[0], "min": 0, "max": 1, "mode": "uniform"},
+    ],
+)
+@pytest.mark.parametrize(
+    "knots_v",
+    [
+        torch.rand(n_knots[1]),
+        {"n": n_knots[1], "min": 0, "max": 1, "mode": "auto"},
+        {"n": n_knots[1], "min": 0, "max": 1, "mode": "uniform"},
+    ],
+)
+@pytest.mark.parametrize(
+    "control_points", [torch.rand(n_ctrl_pts, n_ctrl_pts), None]
+)
+@pytest.mark.parametrize("pts", points)
+def test_forward(knots_u, knots_v, control_points, pts):
+
+    # Define the model
+    model = SplineSurface(
+        orders=orders,
+        knots_u=knots_u,
+        knots_v=knots_v,
+        control_points=control_points,
+    )
+
+    # Evaluate the model
+    output_ = model(pts)
+    assert output_.shape == (*pts.shape[:-1], 1)
+
+
+@pytest.mark.parametrize(
+    "knots_u",
+    [
+        torch.rand(n_knots[0]),
+        {"n": n_knots[0], "min": 0, "max": 1, "mode": "auto"},
+        {"n": n_knots[0], "min": 0, "max": 1, "mode": "uniform"},
+    ],
+)
+@pytest.mark.parametrize(
+    "knots_v",
+    [
+        torch.rand(n_knots[1]),
+        {"n": n_knots[1], "min": 0, "max": 1, "mode": "auto"},
+        {"n": n_knots[1], "min": 0, "max": 1, "mode": "uniform"},
+    ],
+)
+@pytest.mark.parametrize(
+    "control_points", [torch.rand(n_ctrl_pts, n_ctrl_pts), None]
+)
+@pytest.mark.parametrize("pts", points)
+def test_backward(knots_u, knots_v, control_points, pts):
+
+    # Define the model
+    model = SplineSurface(
+        orders=orders,
+        knots_u=knots_u,
+        knots_v=knots_v,
+        control_points=control_points,
+    )
+
+    # Evaluate the model
+    output_ = model(pts)
+    loss = torch.mean(output_)
+    loss.backward()
+    assert model.control_points.grad.shape == model.control_points.shape