add mutual solver-weighting link

tests/test_weighting/test_ntk_weighting.py

@@ -2,64 +2,32 @@ import pytest
 from pina import Trainer
 from pina.solver import PINN
 from pina.model import FeedForward
-from pina.problem.zoo import Poisson2DSquareProblem
 from pina.loss import NeuralTangentKernelWeighting
+from pina.problem.zoo import Poisson2DSquareProblem
 
+
+# Initialize problem and model
 problem = Poisson2DSquareProblem()
-condition_names = problem.conditions.keys()
+problem.discretise_domain(10)
+model = FeedForward(len(problem.input_variables), len(problem.output_variables))
 
 
-@pytest.mark.parametrize(
-    "model,alpha",
-    [
-        (
-            FeedForward(
-                len(problem.input_variables), len(problem.output_variables)
-            ),
-            0.5,
-        )
-    ],
-)
-def test_constructor(model, alpha):
-    NeuralTangentKernelWeighting(model=model, alpha=alpha)
+@pytest.mark.parametrize("alpha", [0.0, 0.5, 1.0])
+def test_constructor(alpha):
+    NeuralTangentKernelWeighting(alpha=alpha)
 
-
-@pytest.mark.parametrize("model", [0.5])
-def test_wrong_constructor1(model):
+    # Should fail if alpha is not >= 0
     with pytest.raises(ValueError):
-        NeuralTangentKernelWeighting(model)
+        NeuralTangentKernelWeighting(alpha=-0.1)
 
-
-@pytest.mark.parametrize(
-    "model,alpha",
-    [
-        (
-            FeedForward(
-                len(problem.input_variables), len(problem.output_variables)
-            ),
-            1.2,
-        )
-    ],
-)
-def test_wrong_constructor2(model, alpha):
+    # Should fail if alpha is not <= 1
     with pytest.raises(ValueError):
-        NeuralTangentKernelWeighting(model, alpha)
+        NeuralTangentKernelWeighting(alpha=1.1)
 
 
-@pytest.mark.parametrize(
-    "model,alpha",
-    [
-        (
-            FeedForward(
-                len(problem.input_variables), len(problem.output_variables)
-            ),
-            0.5,
-        )
-    ],
-)
-def test_train_aggregation(model, alpha):
-    weighting = NeuralTangentKernelWeighting(model=model, alpha=alpha)
-    problem.discretise_domain(50)
+@pytest.mark.parametrize("alpha", [0.0, 0.5, 1.0])
+def test_train_aggregation(alpha):
+    weighting = NeuralTangentKernelWeighting(alpha=alpha)
     solver = PINN(problem=problem, model=model, weighting=weighting)
     trainer = Trainer(solver=solver, max_epochs=5, accelerator="cpu")
     trainer.train()