Browse Source

adding generalized problem class

phillip.rothenbeck 2 years ago
parent
commit
468738d478
1 changed files with 24 additions and 0 deletions
  1. 24 0
      src/problem.py

+ 24 - 0
src/problem.py

@@ -0,0 +1,24 @@
+import torch
+from .dataset import PandemicDataset
+
+class PandemicProblem:
+    def __init__(self, data: PandemicDataset) -> None:
+        """Parent class for all pandemic problem classes. Holding the function, that calculates the residuals of the differential system.
+
+        Args:
+            data (PandemicDataset): Dataset holding the time values used.
+        """
+        #store the gradients for each group
+        self.gradients = [torch.zeros((len(data.t_raw), data.number_groups)) for _ in range(data.number_groups)]
+
+        for i in range(data.number_groups):
+            self.gradients[i][:, i] = 1
+
+    def to_device(self, device):
+        for i in range(len(self.gradients)):
+            self.gradients[i] = self.gradients[i].to(device)
+
+    def residual():
+        """NEEDS TO BE IMPLEMENTED WHEN INHERITING FROM THIS CLASS
+        """
+        pass