diff --git a/gempy/core/data/orientations.py b/gempy/core/data/orientations.py index a8222f924..fdf655a74 100644 --- a/gempy/core/data/orientations.py +++ b/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. diff --git a/gempy/core/data/structural_frame.py b/gempy/core/data/structural_frame.py index 3c2aafb49..968f0305a 100644 --- a/gempy/core/data/structural_frame.py +++ b/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, diff --git a/gempy/core/data/surface_points.py b/gempy/core/data/surface_points.py index bf4ead77f..4a50fd9c7 100644 --- a/gempy/core/data/surface_points.py +++ b/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: diff --git a/test/test_core/test_data/test_orientations.py b/test/test_core/test_data/test_orientations.py new file mode 100644 index 000000000..dae8adba7 --- /dev/null +++ b/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'] \ No newline at end of file diff --git a/test/test_core/test_data/test_surface_points.py b/test/test_core/test_data/test_surface_points.py new file mode 100644 index 000000000..fb315719d --- /dev/null +++ b/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'] \ No newline at end of file