- [ENH] Revive orientations from surface points in a light way

Author: Leguark

gempy_plugins/orientations_generator/__init__.py

@@ -1 +1 @@
-from ._orientations_generator import select_nearest_surfaces_points
+from ._orientations_generator import select_nearest_surfaces_points, NearestSurfacePointsSearcher

gempy_plugins/orientations_generator/_orientations_generator.py

@@ -12,7 +12,7 @@ class NearestSurfacePointsSearcher(enum.Enum):
 
 
 def select_nearest_surfaces_points(surface_points_xyz: np.ndarray, searchcrit: Optional[int|float] = 3,
-                                      search_type: NearestSurfacePointsSearcher = NearestSurfacePointsSearcher.KNN
+                                   search_type: NearestSurfacePointsSearcher = NearestSurfacePointsSearcher.KNN
                                    ) -> np.ndarray:
     match search_type:
         case NearestSurfacePointsSearcher.KNN:
@@ -26,121 +26,4 @@ def select_nearest_surfaces_points(surface_points_xyz: np.ndarray, searchcrit: O
         case _:
             raise ValueError(f"Invalid search type: {search_type}")
 
-    return neighbours_surfaces
-
-
-def _select_nearest_surfaces_points(geo_model, surface_points, searchcrit):
-    """
-    Find the neighbour points of the same surface
-    by given radius (radius-search) or fix number (knn).
-    
-    Parameters
-    ----------
-    geo_model : geo_model
-        GemPy-model.
-    surface_points: Pandas-dataframe
-        Contains the dataframe of the (point-)data from the GemPy-model.
-    searchcrit : int or float
-        if is int: uses knn-search.
-        if is float: uses radius-search.
-    """
-
-    # extract surface names
-    surfaces = np.unique(surface_points['surface'])
-    neighbours = []
-    # for each surface
-    if isinstance(searchcrit, int):  # in case knn-search
-        searchcrit = searchcrit + 1  # because the point itself is also found
-        for s in range(surfaces.size):
-            # extract point-ids
-            i_surfaces = surface_points['surface'] == surfaces[s]
-            # extract point coordinates
-            p_surfaces = surface_points[i_surfaces][['X', 'Y', 'Z']]
-            # create search-tree
-            Tree = NearestNeighbors(n_neighbors=searchcrit)
-            # add data to tree
-            Tree.fit(p_surfaces)
-            # find neighbours
-            neighbours_surfaces = Tree.kneighbors(p_surfaces, n_neighbors=searchcrit,
-                                                  return_distance=False)
-            # add neighbours with initial index to total list
-            for n in neighbours_surfaces:
-                neighbours.append(p_surfaces.index[n])
-    else:  # in case radius-search
-        for s in range(surfaces.size):
-            # extract point-ids
-            i_surfaces = surface_points['surface'] == surfaces[s]
-            # extract point coordinates
-            p_surfaces = surface_points[i_surfaces][['X', 'Y', 'Z']]
-            # create search-tree
-            Tree = NearestNeighbors(radius=searchcrit)
-            # add data to tree
-            Tree.fit(p_surfaces)
-            # find neighbours (attention: relativ index!)
-            neighbours_surfaces = Tree.radius_neighbors(p_surfaces,
-                                                        radius=searchcrit,
-                                                        return_distance=False)
-            # add neighbours with initial index to total list
-            for n in neighbours_surfaces:
-                neighbours.append(p_surfaces.index[n])
-    return neighbours
-
-
-def set_orientation_from_neighbours(geo_model, neighbours):
-    """
-    Calculates the orientation of one point with its neighbour points
-    of the same surface.
-    Parameters
-    ----------
-    geo_model : geo_model
-        GemPy-model.
-    neighbours : Int64Index
-        point-neighbours-id, first id is the point itself.
-    """
-
-    # compute normal vector for the point
-    if neighbours.size > 2:
-        # extract point coordinates
-        coo = geo_model._surface_points.df.loc[neighbours][['X', 'Y', 'Z']]
-        # calculates covariance matrix
-        cov = np.cov(coo.T)
-        # calculate normalized normal vector
-        normvec = normalize(np.cross(cov[0].T, cov[1].T).reshape(1, -1))[0]
-        # check orientation of normal vector (has to be oriented to sky)
-        if normvec[2] < 0:
-            normvec = normvec * (-1)
-        # append to the GemPy-model
-        geo_model.add_orientations(geo_model._surface_points.df['X'][neighbours[0]],
-                                   geo_model._surface_points.df['Y'][neighbours[0]],
-                                   geo_model._surface_points.df['Z'][neighbours[0]],
-                                   geo_model._surface_points.df['surface'][neighbours[0]],
-                                   normvec.tolist())
-    # if computation is impossible set normal vector to default orientation
-    else:
-        print("orientation calculation of point" + str(neighbours[0]) + "is impossible")
-        print("-> default vector is set [0,0,1]")
-        geo_model.add_orientations(geo_model._surface_points.df['X'][neighbours[0]],
-                                   geo_model._surface_points.df['Y'][neighbours[0]],
-                                   geo_model._surface_points.df['Z'][neighbours[0]],
-                                   geo_model._surface_points.df['surface'][neighbours[0]],
-                                   orientation=[0, 0, 1])
-    return geo_model._orientations
-
-
-def set_orientation_from_neighbours_all(geo_model, neighbours):
-    """
-    Calculates the orientations for all points with given neighbours.
-    Parameters
-    ----------
-    geo_model : geo_model
-        GemPy-model.
-    neighbours : list of Int64Index
-        point-neighbours-IDs, the first id is the id of the point
-        for which the orientation is calculated.
-    """
-
-    # compute normal vector for the points
-    for n in neighbours:
-        set_orientation_from_neighbours(geo_model, n)
-
-    return geo_model._orientations
+    return neighbours_surfaces

tests/test_orientations_from_surface_points.py

@@ -1,12 +1,9 @@
 import gempy as gp
-import numpy as np
 import os
 
-from orientations_generator import select_nearest_surfaces_points
-from orientations_generator._orientations_generator import NearestSurfacePointsSearcher
+from orientations_generator import select_nearest_surfaces_points, NearestSurfacePointsSearcher
 
 input_path = os.path.dirname(__file__) + '/../../examples/data'
-
 data_path = os.path.abspath('data')
 
 
@@ -26,7 +23,9 @@ def _model_factory():
 def test_set_orientations():
     geo_data = _model_factory()
 
-    orientations: gp.data.OrientationsTable = gp.create_orientations_from_surface_points(geo_data.surface_points)
+    orientations: gp.data.OrientationsTable = gp.create_orientations_from_surface_points_coords(
+        xyz_coords=geo_data.surface_points
+    )
 
     gp.add_orientations(
         geo_model=geo_data,
@@ -43,16 +42,14 @@ def test_select_nearest_surface_points():
     geo_model = _model_factory()
     print(geo_model)
 
-    # TODO: This is to select the surface points we want to use to calculate the orientations
-    
     element: gp.data.StructuralElement = geo_model.structural_frame.get_element_by_name("fault")
 
     # find neighbours
     knn = select_nearest_surfaces_points(
         surface_points_xyz=element.surface_points.xyz,
         searchcrit=3
     )
-    
+
     radius = select_nearest_surfaces_points(
         surface_points_xyz=element.surface_points.xyz,
         searchcrit=200.,
@@ -63,42 +60,30 @@ def test_select_nearest_surface_points():
 
 
 def test_set_orientation_from_neighbours():
-    
-    
-    
-    
-    # find neighbours
-    neighbours = gp.select_nearest_surfaces_points(geo_data, fault_poi, 5)
-    # calculate one fault orientation
-    gp.set_orientation_from_neighbours(geo_data, neighbours[1])
-    # find the calculated orientation
-    test = geo_data._orientations.df.sort_index().iloc[-1][['dip', 'azimuth']].values
-
-    # calculate reference
-    reference = [90 - np.arctan(0.5) / np.pi * 180, 90]
-
-    assert np.array_equal(reference, test)
-
-
-def test_set_orientation_from_neighbours_all():
-    data_path = 'https://raw.githubusercontent.com/cgre-aachen/gempy_data/master/'
-    path_to_data = data_path + "/data/input_data/jan_models/"
+    geo_model = _model_factory()
 
-    geo_data = gp.create_data_legacy('fault', extent=[0, 1000, 0, 1000, 0, 1000],
-                                     resolution=[50, 50, 50],
-                                     path_o=path_to_data + "model5_orientations.csv",
-                                     path_i=path_to_data + "model5_surface_points.csv")
+    print(geo_model)
 
-    # count orientations before orientation calculation
-    length_pre = geo_data._orientations.df.shape[0]
+    element: gp.data.StructuralElement = geo_model.structural_frame.get_element_by_name("fault")
 
     # find neighbours
-    neighbours = gp.select_nearest_surfaces_points(geo_data, geo_data._surface_points.df, 2)
-    # calculate all fault orientations
-    gp.set_orientation_from_neighbours_all(geo_data, neighbours)
+    knn = select_nearest_surfaces_points(
+        surface_points_xyz=element.surface_points.xyz,
+        searchcrit=3
+    )
 
-    # count orientations after orientation calculation
-    length_after = geo_data._orientations.df.shape[0]
+    orientations: gp.data.OrientationsTable = gp.create_orientations_from_surface_points_coords(
+        xyz_coords=geo_model.surface_points.xyz,
+        subset=knn
+    )       
 
-    assert np.array_equal(geo_data._surface_points.df.shape[0],
-                          length_after - length_pre)
+    gp.add_orientations(
+        geo_model=geo_model,
+        x=orientations.data['X'],
+        y=orientations.data['Y'],
+        z=orientations.data['Z'],
+        pole_vector=orientations.grads,
+        elements_names=geo_model.structural_frame.elements_names[0],
+    )
+    
+    return knn