black code formatter + .rst docs

Author: AleDinve

docs/source/_rst/loss/ntk_weighting.rst

@@ -0,0 +1,9 @@
+NeuralTangentKernelWeighting
+===============
+.. currentmodule:: pina.loss
+
+.. automodule:: pina.loss
+
+.. autoclass:: WeightingInterface
+   :members:
+   :show-inheritance:

pina/loss/__init__.py

@@ -8,12 +8,12 @@
     "PowerLoss",
     "WeightingInterface",
     "ScalarWeighting",
-    "Neural_Tangent_Kernel"
+    "NeuralTangetKernelWeighting",
 ]
 
 from .loss_interface import LossInterface
 from .power_loss import PowerLoss
 from .lp_loss import LpLoss
 from .weighting_interface import WeightingInterface
 from .scalar_weighting import ScalarWeighting
-from .ntk_weighting import Neural_Tangent_Kernel
+from .ntk_weighting import NeuralTangetKernelWeighting

pina/loss/ntk_weighting.py

@@ -1,61 +1,59 @@
-""" Module for Neural Tangent Kernel Class """
+"""Module for Neural Tangent Kernel Class"""
 
+import torch
 from torch.nn import Module
-from torch import norm,cat
 from .weighting_interface import WeightingInterface
 from ..operator import grad
 from ..utils import check_consistency
 
-class Neural_Tangent_Kernel(WeightingInterface):
+
+class NeuralTangetKernelWeighting(WeightingInterface):
     """
     TODO
     """
+
     def __init__(self, model, alpha=0.5):
         super().__init__()
         check_consistency(alpha, float)
-        
-        if alpha<0 or alpha > 1:
+        check_consistency(model, Module)
+        if alpha < 0 or alpha > 1:
             raise ValueError("alpha should be a value between 0 and 1")
-        if not isinstance(model, Module):
-            raise TypeError("Please pass a valid torch.nn.Module model")
         self.alpha = alpha
         self.model = model
         self.weights = {}
         self.default_value_weights = 1
 
     def aggregate(self, losses):
         """
-        Weights the losses according to the Neural Tangent Kernel 
+        Weights the losses according to the Neural Tangent Kernel
 
         :param dict(torch.Tensor) input: The dictionary of losses.
         :param alpha(float) input: The parameter alpha that regulates the moving average
-        between old and new weights. 
+        between old and new weights.
         :return: The losses aggregation. It should be a scalar Tensor.
         :rtype: torch.Tensor
 
-        Reference: TODO
-        Wang, S., Sankaran, S., Wang, H., & Perdikaris, P. (2023). 
-        An expert's guide to training physics-informed neural networks. 
+        Reference:
+        Wang, S., Sankaran, S., Wang, H., & Perdikaris, P. (2023).
+        An expert's guide to training physics-informed neural networks.
         arXiv preprint arXiv:2308.08468.
-        """         
+        """
         losses_norm = {}
         for condition in losses:
             losses[condition].backward(retain_graph=True)
-            grads=[]
+            grads = []
             for param in self.model.parameters():
                 grads.append(param.grad.view(-1))
-            grads = cat(grads)
-            
-            print(grads)
-            losses_norm[condition]=norm(grads)
+            grads = torch.cat(grads)
+            losses_norm[condition] = torch.norm(grads)
         self.weights = {
-            condition: 
-            self.alpha * self.weights.get(
-                condition, self.default_value_weights) + 
-            (1- self.alpha)*losses_norm[condition]/sum(losses_norm.values())
+            condition: self.alpha
+            * self.weights.get(condition, self.default_value_weights)
+            + (1 - self.alpha)
+            * losses_norm[condition]
+            / torch.sum(losses_norm.values())
             for condition in losses
         }
         return sum(
-            self.weights.get(condition, self.default_value_weights) * loss for
-            condition, loss in losses.items()
-        )
+            self.weights[condition] * loss for condition, loss in losses.items()
+        )

tests/test_weighting/test_ntk_weighting.py

@@ -3,28 +3,31 @@
 from pina.solver import PINN
 from pina.model import FeedForward
 from pina.problem.zoo import Poisson2DSquareProblem
-from pina.loss import Neural_Tangent_Kernel
+from pina.loss import NeuralTangetKernelWeighting
 
 problem = Poisson2DSquareProblem()
 model = FeedForward(len(problem.input_variables), len(problem.output_variables))
 condition_names = problem.conditions.keys()
-#print(problem.conditions.keys())
+# print(problem.conditions.keys())
+
+
+def test_constructor(model, alpha):
+    NeuralTangetKernelWeighting(model=model, alpha=alpha)
 
-def test_constructor(model,alpha):
-    Neural_Tangent_Kernel(model=model,alpha=alpha)
 
 def test_wrong_constructor1(alpha):
     with pytest.raises(ValueError):
-        Neural_Tangent_Kernel(model,alpha)
+        NeuralTangetKernelWeighting(model, alpha)
+
 
 def test_wrong_constructor2(model):
     with pytest.raises(TypeError):
-        Neural_Tangent_Kernel(model)
+        NeuralTangetKernelWeighting(model)
+
 
 def test_train_aggregation(model):
-    weighting = Neural_Tangent_Kernel(model=model,alpha = 0.5)
+    weighting = NeuralTangetKernelWeighting(model=model, alpha=0.5)
     problem.discretise_domain(50)
     solver = PINN(problem=problem, model=model, weighting=weighting)
     trainer = Trainer(solver=solver, max_epochs=5, accelerator="cpu")
     trainer.train()
-