[ENH] Mesh extraction on pytorch

Author: Leguark

gempy_engine/API/dual_contouring/_dual_contouring.py

@@ -3,6 +3,9 @@
 
 import numpy as np
 
+from gempy_engine.config import AvailableBackends
+
+from ...core.backend_tensor import BackendTensor
 from ...core.data.dual_contouring_data import DualContouringData
 from ...core.data.dual_contouring_mesh import DualContouringMesh
 from ...core.utils import gempy_profiler_decorator
@@ -49,12 +52,12 @@ def compute_dual_contouring(dc_data_per_stack: DualContouringData, left_right_co
             # * Fancy triangulation 👗
 
             # * Average gradient for the edges
-            from gempy_engine.core.backend_tensor import BackendTensor
             edges_normals = BackendTensor.t.zeros((valid_edges.shape[0], 12, 3), dtype=BackendTensor.dtype_obj)
             edges_normals[:] = np.nan
             edges_normals[valid_edges] = dc_data_per_stack.gradients[slice_object]
 
-            if LEGACY:=False:
+            # if LEGACY:=True:
+            if BackendTensor.engine_backend != AvailableBackends.PYTORCH:
                 with warnings.catch_warnings():
                     warnings.simplefilter("ignore", category=RuntimeWarning)
                     voxel_normal  = np.nanmean(edges_normals, axis=1)
@@ -76,7 +79,7 @@ def compute_dual_contouring(dc_data_per_stack: DualContouringData, left_right_co
                 voxel_normal = sum_normals / valid_count.clamp(min=1)
 
                 # Remove rows where all elements were NaN (and hence valid_count is 0)
-                voxel_normal = voxel_normal[valid_count > 0]
+                voxel_normal = voxel_normal[valid_count > 0].reshape(-1, 3)
 
 
             valid_voxels = dc_data_per_surface.valid_voxels
@@ -89,5 +92,11 @@ def compute_dual_contouring(dc_data_per_stack: DualContouringData, left_right_co
             indices = np.vstack(indices)
 
         # @on
-        stack_meshes.append(DualContouringMesh(vertices, indices, dc_data_per_stack))
+        stack_meshes.append(
+            DualContouringMesh(
+                BackendTensor.t.to_numpy(vertices),
+                indices,
+                dc_data_per_stack
+            )
+        )
     return stack_meshes

gempy_engine/modules/dual_contouring/dual_contouring_interface.py

@@ -173,7 +173,8 @@ def generate_dual_contouring_vertices(dc_data_per_stack: DualContouringData, sli
         bias_xyz[isclose] = np.nan  # np zero values to nans
         mass_points = np.nanmean(bias_xyz, axis=1)  # Mean ignoring nans
     else:  # ? This is actually doing something
-        bias_xyz = BackendTensor.t.copy(edges_xyz[:, :12]).detach().numpy()
+        bias_xyz = BackendTensor.t.copy(edges_xyz[:, :12])
+        bias_xyz = BackendTensor.t.to_numpy(bias_xyz)
         mask = bias_xyz == 0
         masked_arr = np.ma.masked_array(bias_xyz, mask)
         mass_points = masked_arr.mean(axis=1)
@@ -200,10 +201,15 @@ def generate_dual_contouring_vertices(dc_data_per_stack: DualContouringData, sli
     # Compute LSTSQS in all voxels at the same time
     A = edges_normals
     b = (A * edges_xyz).sum(axis=2)
-    
-    term1 = BackendTensor.t.einsum("ijk, ilj->ikl", A, BackendTensor.t.transpose(A, (2, 1)))
+
+    if BackendTensor.engine_backend == AvailableBackends.PYTORCH:
+        transpose_shape = (2, 1)
+    else:
+        transpose_shape = (0, 2,1)
+        
+    term1 = BackendTensor.t.einsum("ijk, ilj->ikl", A, BackendTensor.t.transpose(A, transpose_shape))
     term2 = BackendTensor.t.linalg.inv(term1)
-    term3 = BackendTensor.t.einsum("ijk,ik->ij", BackendTensor.t.transpose(A, (2, 1)), b)
+    term3 = BackendTensor.t.einsum("ijk,ik->ij", BackendTensor.t.transpose(A, transpose_shape), b)
     vertices = BackendTensor.t.einsum("ijk, ij->ik", term2, term3)
 
     if debug:

gempy_engine/modules/dual_contouring/fancy_triangulation.py

@@ -178,9 +178,10 @@ def check_voxels_exist_next_to_edge(coord_col, edge_vector, _left_right_array_ac
 
     valid_edges_within_extent = code__a_prod_edge * code__b_prod_edge * code__c_prod_edge  # * Valid in the sense that there are valid voxels around
 
-    code__a_p = mapped_voxel_0[:, valid_edges_within_extent] == 0  # (n_voxels, n_voxels - active_voxels_for_given_edge - invalid_edges - edges_at_extent_border)
-    code__b_p = mapped_voxel_1[:, valid_edges_within_extent] == 0  # (n_voxels, n_voxels - active_voxels_for_given_edge - invalid_edges - edges_at_extent_border)
-    code__c_p = mapped_voxel_2[:, valid_edges_within_extent] == 0  # (n_voxels, n_voxels - active_voxels_for_given_edge - invalid_edges - edges_at_extent_border)
+    from ...core.backend_tensor import BackendTensor
+    code__a_p = BackendTensor.t.array(mapped_voxel_0[:, valid_edges_within_extent] == 0)  # (n_voxels, n_voxels - active_voxels_for_given_edge - invalid_edges - edges_at_extent_border)
+    code__b_p = BackendTensor.t.array(mapped_voxel_1[:, valid_edges_within_extent] == 0)  # (n_voxels, n_voxels - active_voxels_for_given_edge - invalid_edges - edges_at_extent_border)
+    code__c_p = BackendTensor.t.array(mapped_voxel_2[:, valid_edges_within_extent] == 0)  # (n_voxels, n_voxels - active_voxels_for_given_edge - invalid_edges - edges_at_extent_border)
 
     if False:
         debug_code_p = code__a_p + code__b_p + code__c_p  # (n_voxels, n_voxels - active_voxels_for_given_edge - invalid_edges - edges_at_extent_border)