feat: data operations

FoldOptLib/datatypes/enums.py

@@ -2,10 +2,9 @@
 
 
 class DataType(IntEnum):
-    FOLDED_AXIAL_SURFACE = 0
-    FOLDED_FOLIATION = 1
-    GEOLOGICAL_KNOWLEDGE = 2
-    BOUNDING_BOX = 3
+    DATA = 0
+    GEOLOGICAL_KNOWLEDGE = 1
+    BOUNDING_BOX = 2
 
 
 class KnowledgeType(IntEnum):

FoldOptLib/input/data_storage.py

@@ -3,16 +3,37 @@
 from ..objective_functions import GeologicalKnowledgeFunctions
 from LoopStructural import BoundingBox
 import pandas
-
+import numpy
+import beartype
 
 @dataclass
 class InputData:
-    folded_axial_surface: pandas.DataFrame
-    folded_foliation: pandas.DataFrame
+    data: pandas.DataFrame
     bounding_box: BoundingBox
     geological_knowledge: InputGeologicalKnowledge = None
 
 
     def __post_init__(self):
-    self.
+        pass 
+
+    def foliations(self): 
+        return numpy.unique(self.data['feature_name'])
+
+    def number_of_foliations(self): 
+        return len(self.foliations())
+
+    def get_foliations(self, feature_name): 
+        return self.data[self.data['feature_name'] == feature_name]
+
+
+    def __getitem__(self, data_type: DataType): 
+
+        data_map = {
+            DataType.DATA: self.data,
+            DataType.GEOLOGICAL_KNOWLEDGE: self.geological_knowledge,
+            DataType.BOUNDING_BOX: self.bounding_box
+        }
+
+
+        return data_map[data_type]