minor fix

Author: ndem0

pina/_src/model/block/kan_block.py

@@ -3,11 +3,16 @@
 import numpy as np
 
 from pina._src.model.spline import Spline
+from pina._src.model.vectorized_spline import VectorizedSpline
 
 
 class KANBlock(torch.nn.Module):
     """define a KAN layer using splines"""
-    def __init__(self, k, input_dimensions, output_dimensions, inner_nodes, num=3, grid_eps=0.1, grid_range=[-1, 1], grid_extension=True, noise_scale=0.1, base_function=torch.nn.SiLU(), scale_base_mu=0.0, scale_base_sigma=1.0, scale_sp=1.0, sparse_init=True, sp_trainable=True, sb_trainable=True):
+    def __init__(self, k, input_dimensions, output_dimensions, inner_nodes,
+    num=3, grid_eps=0.1, grid_range=[-1, 1], grid_extension=True,
+    noise_scale=0.1, base_function=torch.nn.SiLU(), scale_base_mu=0.0,
+    scale_base_sigma=1.0, scale_sp=1.0, sparse_init=True, sp_trainable=True,
+    sb_trainable=True):
         """
         Initialize the KAN layer.
 
@@ -46,7 +51,6 @@ def __init__(self, k, input_dimensions, output_dimensions, inner_nodes, num=3, g
         # )
         # print(control_points.shape)
         if self.vec:
-            from pina.model.spline import SplineVectorized as VectorizedSpline
             control_points = torch.randn(self.input_dimensions * self.output_dimensions, n_control_points)
             print('control points', control_points.shape)
             control_points = torch.stack([

tests/test_model/test_kolmogorov_arnold_network.py

@@ -1,7 +1,6 @@
 import torch
 import pytest
 
-from pina.model.block import KANBlock
 from pina.model import KolmogorovArnoldNetwork
 
 data = torch.rand((20, 3))
@@ -81,16 +80,14 @@ def test_train_poisson():
     solver = PINN(model=model, problem=problem)
     trainer = Trainer(
         solver=solver,
-        max_epochs=1000,
+        max_epochs=10,
         accelerator="cpu",
         batch_size=100,
         train_size=1.0,
         val_size=0.0,
         test_size=0.0,
     )
     trainer.train()
-    assert False
-
 
 
 

tests/test_model/test_spline.py

@@ -193,22 +193,22 @@ def test_derivative(args, pts):
     assert torch.allclose(first_der, first_der_auto, atol=1e-4, rtol=1e-4)
 
 
-@pytest.mark.parametrize("out_dim", [1, 3, 5])
-def test_vectorized(out_dim):
+#@pytest.mark.parametrize("args", valid_args) # TODO
+def test_vectorized():
 
     N = 7
     cps = []
     splines = []
     for i in range(N):
-        cp = torch.rand(n_ctrl_pts, 3)
+        cp = torch.rand(n_ctrl_pts)
         cps.append(cp)
         spline = Spline(
             order=order,
             control_points=cp
         )
         splines.append(spline)
 
-    from pina.model.spline import SplineVectorized as VectorizedSpline
+    from pina.model import VectorizedSpline
     unique_cps = torch.stack(cps, dim=0)
     print(unique_cps.shape)
     print(cps[0].shape)
@@ -224,7 +224,8 @@ def test_vectorized(out_dim):
     result_single = torch.stack([
         splines[i](x) for i in range(N)
     ])
-    result_single = result_single.squeeze(2).permute(1, 0, 2)
+    print(result_single.shape)
+    result_single = result_single.permute(1, 2, 0)
     out_vectorized = vectorized_spline(x)
     print(out_vectorized.shape)
     print(result_single.shape)