Mapping formations to surface_points_copy and orientations DataFrames

gempy/core/data/orientations.py

@@ -21,11 +21,11 @@ class OrientationsTable:
 
     """
 
-    dt = np.dtype([('X', 'f8'), ('Y', 'f8'), ('Z', 'f8'), ('G_x', 'f8'), ('G_y', 'f8'), ('G_z', 'f8'), ('id', 'i4'), ('nugget', 'f8')])  #: The custom data type for the data array.
+    dt = np.dtype([('X', 'f8'), ('Y', 'f8'), ('Z', 'f8'), ('G_x', 'f8'), ('G_y', 'f8'), ('G_z', 'f8'), ('id', 'i4'), ('nugget', 'f8'), ('formation', 'U20')])  #: The custom data type for the data array.
     data: np.ndarray = Field(
         default=np.zeros(0, dtype=dt),
         exclude=True,
-        description="A structured NumPy array holding the X, Y, Z coordinates, gradients G_x, G_y, G_z, id, and nugget of each orientation.",
+        description="A structured NumPy array holding the X, Y, Z coordinates, gradients G_x, G_y, G_z, id, nugget and formation of each orientation.",
     )  #: A structured NumPy array holding the X, Y, Z coordinates, id, and nugget of each surface point.
     name_id_map: Optional[dict[str, int]] = None  #: A mapping between orientation names and ids.
 
@@ -73,7 +73,7 @@ def _data_from_arrays(cls, x, y, z, G_x, G_y, G_z, names, nugget, name_id_map=No
         else:
             ids = np.array([name_id_map[name] for name in names])
         data = np.zeros(len(x), dtype=OrientationsTable.dt)
-        data['X'], data['Y'], data['Z'], data['G_x'], data['G_y'], data['G_z'], data['id'], data['nugget'] = x, y, z, G_x, G_y, G_z, ids, nugget
+        data['X'], data['Y'], data['Z'], data['G_x'], data['G_y'], data['G_z'], data['id'], data['nugget'], data['formation'] = x, y, z, G_x, G_y, G_z, ids, nugget, names
         return data, name_id_map
 
     @classmethod
@@ -143,6 +143,10 @@ def nugget(self) -> np.ndarray:
         """
         return self.data['nugget']
 
+    @property
+    def formation(self) -> np.ndarray:
+        return self.data['formation']
+
     @property
     def ids(self) -> np.ndarray:
         """Get the IDs.

gempy/core/data/structural_frame.py

@@ -59,7 +59,7 @@ def from_data_tables(cls, surface_points: SurfacePointsTable, orientations: Orie
                 orientation_i = OrientationsTable.empty_orientation(id_)
 
             structural_element: StructuralElement = StructuralElement(
-                name=surface_points.id_to_name(i),
+                name=surface_points.id_to_name(surface_points_groups[i].df['id'].unique()),
                 id=id_,
                 surface_points=surface_points_groups[i],
                 orientations=orientation_i,

gempy/core/data/surface_points.py

@@ -22,12 +22,12 @@ class SurfacePointsTable:
     A dataclass to represent a table of surface points in a geological model.
 
     """
-    dt = np.dtype([('X', 'f8'), ('Y', 'f8'), ('Z', 'f8'), ('id', 'i4'), ('nugget', 'f8')])  #: The custom data type for the data array.
+    dt = np.dtype([('X', 'f8'), ('Y', 'f8'), ('Z', 'f8'), ('id', 'i4'), ('nugget', 'f8'), ('formation', 'U20')])  #: The custom data type for the data array.
 
     data: np.ndarray = Field(
         default=np.zeros(0, dtype=dt),
         exclude=True,
-        description="A structured NumPy array holding the X, Y, Z coordinates, id, and nugget of each surface point."
+        description="A structured NumPy array holding the X, Y, Z coordinates, id, nugget and formation of each surface point."
     )  #: A structured NumPy array holding the X, Y, Z coordinates, id, and nugget of each surface point.
     name_id_map: Optional[dict[str, int]] = None  #: A mapping between surface point names and ids.
     _model_transform: Optional[Transform] = None
@@ -94,7 +94,7 @@ def _data_from_arrays(cls, x: np.ndarray, y: np.ndarray, z: np.ndarray,
             ids = np.array([name_id_map[name] for name in names])
 
         data = np.zeros(len(x), dtype=SurfacePointsTable.dt)
-        data['X'], data['Y'], data['Z'], data['id'], data['nugget'] = x, y, z, ids, nugget
+        data['X'], data['Y'], data['Z'], data['id'], data['nugget'], data['formation'] = x, y, z, ids, nugget, names
         return data, name_id_map
 
     @classmethod
@@ -115,7 +115,7 @@ def id_to_name(self, id: int) -> str:
         Returns:
             str: The name of the surface point.
         """
-        return list(self.name_id_map.keys())[id]
+        return [key for key, value in self.name_id_map.items() if value == 	id][0]
 
     @property
     def xyz(self) -> np.ndarray:
@@ -137,6 +137,10 @@ def nugget(self) -> np.ndarray:
     def nugget(self, value: np.ndarray):
         self.data['nugget'] = value
 
+    @property
+    def formation(self) -> np.ndarray:
+        return self.data['formation']
+
     @property
     def model_transform(self) -> Transform:
         if self._model_transform is None:

test/test_core/test_data/test_orientations.py

@@ -0,0 +1,7 @@
+from gempy.core.data import GeoModel
+from test.test_api.test_initialization_and_compute_api import _create_data
+
+def test_surface_points_copy_df():
+    geo_data: GeoModel = _create_data()
+
+    assert list(geo_data.structural_frame.surface_points_copy.df.columns) == ['X', 'Y', 'Z', 'id', 'nugget', 'formation']

test/test_core/test_data/test_surface_points.py

@@ -0,0 +1,17 @@
+from gempy.core.data import GeoModel
+from test.test_api.test_initialization_and_compute_api import _create_data
+
+def test_id_to_name():
+    geo_data: GeoModel = _create_data()
+
+    id_number1 = geo_data.structural_frame.surface_points_copy.df['id'].unique()[0]
+    id_number2 = geo_data.structural_frame.surface_points_copy.df['id'].unique()[1]
+
+    assert geo_data.structural_frame.surface_points_copy.id_to_name(id_number1) == 'rock1'
+    assert geo_data.structural_frame.surface_points_copy.id_to_name(id_number2) == 'rock2'
+
+
+def test_surface_points_copy_df():
+    geo_data: GeoModel = _create_data()
+
+    assert list(geo_data.structural_frame.surface_points_copy.df.columns) == ['X', 'Y', 'Z', 'id', 'nugget', 'formation']