optimizer and scheduler classes

pina/optimizer.py

@@ -0,0 +1,21 @@
+""" Module for PINA Optimizer """
+
+import torch
+from .utils import check_consistency
+
+class Optimizer: # TODO improve interface
+    pass
+
+
+class TorchOptimizer(Optimizer):
+
+    def __init__(self, optimizer_class, **kwargs):
+        check_consistency(optimizer_class, torch.optim.Optimizer, subclass=True)
+
+        self.optimizer_class = optimizer_class
+        self.kwargs = kwargs
+
+    def hook(self, parameters):
+        self.optimizer_instance = self.optimizer_class(
+            parameters, **self.kwargs
+        )