weighting refactory

Co-authored-by: Dario Coscia <dariocos99@gmail.com>

tests/test_weighting/test_ntk_weighting.py

@@ -12,22 +12,42 @@ problem.discretise_domain(10)
 model = FeedForward(len(problem.input_variables), len(problem.output_variables))
 
 
+@pytest.mark.parametrize("update_every_n_epochs", [1, 10, 100, 1000])
 @pytest.mark.parametrize("alpha", [0.0, 0.5, 1.0])
-def test_constructor(alpha):
-    NeuralTangentKernelWeighting(alpha=alpha)
+def test_constructor(update_every_n_epochs, alpha):
+    NeuralTangentKernelWeighting(
+        update_every_n_epochs=update_every_n_epochs, alpha=alpha
+    )
 
     # Should fail if alpha is not >= 0
     with pytest.raises(ValueError):
-        NeuralTangentKernelWeighting(alpha=-0.1)
+        NeuralTangentKernelWeighting(
+            update_every_n_epochs=update_every_n_epochs, alpha=-0.1
+        )
 
     # Should fail if alpha is not <= 1
     with pytest.raises(ValueError):
         NeuralTangentKernelWeighting(alpha=1.1)
 
+    # Should fail if update_every_n_epochs is not an integer
+    with pytest.raises(AssertionError):
+        NeuralTangentKernelWeighting(update_every_n_epochs=1.5)
 
+    # Should fail if update_every_n_epochs is not > 0
+    with pytest.raises(AssertionError):
+        NeuralTangentKernelWeighting(update_every_n_epochs=0)
+
+    # Should fail if update_every_n_epochs is not > 0
+    with pytest.raises(AssertionError):
+        NeuralTangentKernelWeighting(update_every_n_epochs=-3)
+
+
+@pytest.mark.parametrize("update_every_n_epochs", [1, 3])
 @pytest.mark.parametrize("alpha", [0.0, 0.5, 1.0])
-def test_train_aggregation(alpha):
-    weighting = NeuralTangentKernelWeighting(alpha=alpha)
+def test_train_aggregation(update_every_n_epochs, alpha):
+    weighting = NeuralTangentKernelWeighting(
+        update_every_n_epochs=update_every_n_epochs, alpha=alpha
+    )
     solver = PINN(problem=problem, model=model, weighting=weighting)
     trainer = Trainer(solver=solver, max_epochs=5, accelerator="cpu")
     trainer.train()

tests/test_weighting/test_scalar_weighting.py

@@ -29,20 +29,6 @@ def test_constructor(weights):
         ScalarWeighting(weights=[1, 2, 3])
 
 
-@pytest.mark.parametrize(
-    "weights", [1, 1.0, dict(zip(condition_names, [1] * len(condition_names)))]
-)
-def test_aggregate(weights):
-    weighting = ScalarWeighting(weights=weights)
-    losses = dict(
-        zip(
-            condition_names,
-            [torch.randn(1) for _ in range(len(condition_names))],
-        )
-    )
-    weighting.aggregate(losses=losses)
-
-
 @pytest.mark.parametrize(
     "weights", [1, 1.0, dict(zip(condition_names, [1] * len(condition_names)))]
 )

tests/test_weighting/test_self_adaptive_weighting.py

@@ -12,26 +12,28 @@ problem.discretise_domain(10)
 model = FeedForward(len(problem.input_variables), len(problem.output_variables))
 
 
-@pytest.mark.parametrize("k", [10, 100, 1000])
-def test_constructor(k):
-    SelfAdaptiveWeighting(k=k)
+@pytest.mark.parametrize("update_every_n_epochs", [10, 100, 1000])
+def test_constructor(update_every_n_epochs):
+    SelfAdaptiveWeighting(update_every_n_epochs=update_every_n_epochs)
 
-    # Should fail if k is not an integer
+    # Should fail if update_every_n_epochs is not an integer
     with pytest.raises(AssertionError):
-        SelfAdaptiveWeighting(k=1.5)
+        SelfAdaptiveWeighting(update_every_n_epochs=1.5)
 
-    # Should fail if k is not > 0
+    # Should fail if update_every_n_epochs is not > 0
     with pytest.raises(AssertionError):
-        SelfAdaptiveWeighting(k=0)
+        SelfAdaptiveWeighting(update_every_n_epochs=0)
 
-    # Should fail if k is not > 0
+    # Should fail if update_every_n_epochs is not > 0
     with pytest.raises(AssertionError):
-        SelfAdaptiveWeighting(k=-3)
+        SelfAdaptiveWeighting(update_every_n_epochs=-3)
 
 
-@pytest.mark.parametrize("k", [2, 3])
-def test_train_aggregation(k):
-    weighting = SelfAdaptiveWeighting(k=k)
+@pytest.mark.parametrize("update_every_n_epochs", [1, 3])
+def test_train_aggregation(update_every_n_epochs):
+    weighting = SelfAdaptiveWeighting(
+        update_every_n_epochs=update_every_n_epochs
+    )
     solver = PINN(problem=problem, model=model, weighting=weighting)
     trainer = Trainer(solver=solver, max_epochs=5, accelerator="cpu")
     trainer.train()