[ENH] Adapt code to do triangulation in GPU

Author: Leguark

gempy_engine/API/dual_contouring/_dual_contouring.py

@@ -1,3 +1,4 @@
+import numpy
 import warnings
 from typing import List
 
@@ -90,18 +91,19 @@ def compute_dual_contouring(dc_data_per_stack: DualContouringData, left_right_co
                 tree_depth       = dc_data_per_surface.tree_depth,
                 voxel_normals     = voxel_normal 
             )
-            indices = np.vstack(indices)
+            indices = BackendTensor.t.concatenate(indices, axis=0)
 
         # @on
         vertices_numpy = BackendTensor.t.to_numpy(vertices)
+        indices_numpy = BackendTensor.t.to_numpy(indices)
 
         if TRIMESH_LAST_PASS := True:
-            vertices_numpy, indices = _last_pass(vertices_numpy, indices)
+            vertices_numpy, indices_numpy = _last_pass(vertices_numpy, indices_numpy)
 
         stack_meshes.append(
             DualContouringMesh(
                 vertices_numpy,
-                indices,
+                indices_numpy,
                 dc_data_per_stack
             )
         )

gempy_engine/core/backend_tensor.py

@@ -184,24 +184,6 @@ def _concatenate(tensors, axis=0, dtype=None):
         def _transpose(tensor, axes=None):
             return tensor.transpose(axes[0], axes[1])
 
-        def _packbits(tensor, axis=None, bitorder="big"):
-            # PyTorch doesn't have packbits, need to implement or use alternative
-            # This is a simplified version - you might need a more complete implementation
-            if bitorder == "little":
-                # Reverse bits for little endian
-                # noinspection PyArgumentList
-                tensor = flip(tensor, axis=[axis] if axis is not None else [-1])
-
-            # Convert boolean to integers and pack
-            tensor_int = tensor.to(torch.uint8)
-            if axis == 0:
-                # Pack along first axis
-                result = torch.zeros(tensor.shape[1], dtype=torch.uint8, device=tensor.device)
-                for i in range(tensor.shape[0]):
-                    result += tensor_int[i] * (2 ** i)
-                return result
-            else:
-                raise NotImplementedError("packbits only implemented for axis=0")
 
         def _packbits(tensor, axis=None, bitorder="big"):
             """
@@ -270,14 +252,27 @@ def _packbits(tensor, axis=None, bitorder="big"):
             else:
                 raise NotImplementedError(f"packbits not implemented for axis={axis}")
 
+
+        def _to_numpy(tensor):
+            """Convert tensor to numpy array, handling GPU tensors properly"""
+            if hasattr(tensor, 'device') and tensor.device.type == 'cuda':
+                # Move to CPU first, then detach and convert to numpy
+                return tensor.cpu().detach().numpy()
+            elif hasattr(tensor, 'detach'):
+                # CPU tensor, just detach and convert
+                return tensor.detach().numpy()
+            else:
+                # Not a torch tensor, return as-is
+                return tensor
+
         cls.tfnp.sum = _sum
         cls.tfnp.repeat = _repeat
         cls.tfnp.expand_dims = lambda tensor, axis: tensor
         cls.tfnp.invert = lambda tensor: ~tensor
         cls.tfnp.flip = lambda tensor, axis: tensor.flip(axis)
         cls.tfnp.hstack = lambda tensors: torch.concat(tensors, dim=1)
         cls.tfnp.array = _array
-        cls.tfnp.to_numpy = lambda tensor: tensor.detach().numpy()
+        cls.tfnp.to_numpy = _to_numpy
         cls.tfnp.min = lambda tensor, axis: tensor.min(axis=axis)[0]
         cls.tfnp.max = lambda tensor, axis: tensor.max(axis=axis)[0]
         cls.tfnp.rint = lambda tensor: tensor.round().type(torch.int32)

gempy_engine/core/data/interp_output.py

@@ -154,6 +154,7 @@ def get_block_from_value_type(self, value_type: ValueType, slice_: slice):
 
         match (BackendTensor.engine_backend):
             case AvailableBackends.PYTORCH:
-                block = block.detach().numpy()
+                block = BackendTensor.t.to_numpy(block)
+                # block = block.to_numpy()
 
         return block[slice_]

gempy_engine/modules/dual_contouring/dual_contouring_interface.py

@@ -1,13 +1,13 @@
+import torch
 from typing import Tuple
 
 from gempy_engine.config import AvailableBackends
 from ...core.backend_tensor import BackendTensor
 from ...core.data.dual_contouring_data import DualContouringData
-import numpy as np
 
 
-def find_intersection_on_edge(_xyz_corners: np.ndarray, scalar_field_on_corners: np.ndarray,
-                              scalar_at_sp: np.ndarray, masking=None) -> Tuple[np.ndarray, np.ndarray]:
+def find_intersection_on_edge(_xyz_corners, scalar_field_on_corners,
+                              scalar_at_sp, masking=None) -> Tuple:
     """This function finds all the intersections for multiple layers per series
     
     - The shape of valid edges is n_surfaces * xyz_corners. Where xyz_corners is 8 * the octree leaf
@@ -27,13 +27,13 @@ def find_intersection_on_edge(_xyz_corners: np.ndarray, scalar_field_on_corners:
     scalar_at_sp = scalar_at_sp.reshape((-1, 1, 1))
 
     n_isosurface = scalar_at_sp.shape[0]
-    xyz_8 = BackendTensor.t.tile(xyz_8, (n_isosurface, 1, 1))  # TODO: Generalize
+    xyz_8 = BackendTensor.tfnp.tile(xyz_8, (n_isosurface, 1, 1))  # TODO: Generalize
 
     # Compute distance of scalar field on the corners
-    scalar_dx = scalar_8[:, :, :4] - scalar_8[:, :, 4:] 
+    scalar_dx = scalar_8[:, :, :4] - scalar_8[:, :, 4:]
     scalar_d_y = scalar_8[:, :, [0, 1, 4, 5]] - scalar_8[:, :, [2, 3, 6, 7]]
     scalar_d_z = scalar_8[:, :, ::2] - scalar_8[:, :, 1::2]
-    
+
     # Add a tiny value to avoid division by zero
     scalar_dx += 1e-10
     scalar_d_y += 1e-10
@@ -53,16 +53,16 @@ def find_intersection_on_edge(_xyz_corners: np.ndarray, scalar_field_on_corners:
     intersect_dx = d_x[:, :, :] * weight_x[:, :, :]
     intersect_dy = d_y[:, :, :] * weight_y[:, :, :]
     intersect_dz = d_z[:, :, :] * weight_z[:, :, :]
-    
+
     # Mask invalid edges
-    valid_edge_x = np.logical_and(weight_x > -0.01, weight_x < 1.01)
-    valid_edge_y = np.logical_and(weight_y > -0.01, weight_y < 1.01)
-    valid_edge_z = np.logical_and(weight_z > -0.01, weight_z < 1.01)
+    valid_edge_x = BackendTensor.tfnp.logical_and(weight_x > -0.01, weight_x < 1.01)
+    valid_edge_y = BackendTensor.tfnp.logical_and(weight_y > -0.01, weight_y < 1.01)
+    valid_edge_z = BackendTensor.tfnp.logical_and(weight_z > -0.01, weight_z < 1.01)
 
     # * Note(miguel) From this point on the arrays become sparse
-    xyz_8_edges = BackendTensor.t.hstack([xyz_8[:, 4:], xyz_8[:, [2, 3, 6, 7]], xyz_8[:, 1::2]])
-    intersect_segment = BackendTensor.t.hstack([intersect_dx, intersect_dy, intersect_dz])
-    valid_edges = BackendTensor.t.hstack([valid_edge_x, valid_edge_y, valid_edge_z])[:, :, 0]
+    xyz_8_edges = BackendTensor.tfnp.hstack([xyz_8[:, 4:], xyz_8[:, [2, 3, 6, 7]], xyz_8[:, 1::2]])
+    intersect_segment = BackendTensor.tfnp.hstack([intersect_dx, intersect_dy, intersect_dz])
+    valid_edges = BackendTensor.tfnp.hstack([valid_edge_x, valid_edge_y, valid_edge_z])[:, :, 0]
     valid_edges = valid_edges > 0
 
     intersection_xyz = xyz_8_edges[valid_edges] + intersect_segment[valid_edges]
@@ -88,13 +88,13 @@ def triangulate_dual_contouring(dc_data_per_surface: DualContouringData):
     x_2 = centers_xyz[valid_voxels][None, :, :]
 
     manhattan = x_1 - x_2
-    zeros = np.isclose(manhattan[:, :, :], 0, .00001)
-    x_direction_neighbour = np.isclose(manhattan[:, :, 0], dx, .00001)
-    nx_direction_neighbour = np.isclose(manhattan[:, :, 0], -dx, .00001)
-    y_direction_neighbour = np.isclose(manhattan[:, :, 1], dy, .00001)
-    ny_direction_neighbour = np.isclose(manhattan[:, :, 1], -dy, .00001)
-    z_direction_neighbour = np.isclose(manhattan[:, :, 2], dz, .00001)
-    nz_direction_neighbour = np.isclose(manhattan[:, :, 2], -dz, .00001)
+    zeros = BackendTensor.tfnp.isclose(manhattan[:, :, :], 0, .00001)
+    x_direction_neighbour = BackendTensor.tfnp.isclose(manhattan[:, :, 0], dx, .00001)
+    nx_direction_neighbour = BackendTensor.tfnp.isclose(manhattan[:, :, 0], -dx, .00001)
+    y_direction_neighbour = BackendTensor.tfnp.isclose(manhattan[:, :, 1], dy, .00001)
+    ny_direction_neighbour = BackendTensor.tfnp.isclose(manhattan[:, :, 1], -dy, .00001)
+    z_direction_neighbour = BackendTensor.tfnp.isclose(manhattan[:, :, 2], dz, .00001)
+    nz_direction_neighbour = BackendTensor.tfnp.isclose(manhattan[:, :, 2], -dz, .00001)
 
     x_direction = x_direction_neighbour * zeros[:, :, 1] * zeros[:, :, 2]
     nx_direction = nx_direction_neighbour * zeros[:, :, 1] * zeros[:, :, 2]
@@ -103,12 +103,12 @@ def triangulate_dual_contouring(dc_data_per_surface: DualContouringData):
     z_direction = z_direction_neighbour * zeros[:, :, 0] * zeros[:, :, 1]
     nz_direction = nz_direction_neighbour * zeros[:, :, 0] * zeros[:, :, 1]
 
-    np.fill_diagonal(x_direction, True)
-    np.fill_diagonal(nx_direction, True)
-    np.fill_diagonal(y_direction, True)
-    np.fill_diagonal(nx_direction, True)
-    np.fill_diagonal(z_direction, True)
-    np.fill_diagonal(nz_direction, True)
+    BackendTensor.tfnp.fill_diagonal(x_direction, True)
+    BackendTensor.tfnp.fill_diagonal(nx_direction, True)
+    BackendTensor.tfnp.fill_diagonal(y_direction, True)
+    BackendTensor.tfnp.fill_diagonal(ny_direction, True)
+    BackendTensor.tfnp.fill_diagonal(z_direction, True)
+    BackendTensor.tfnp.fill_diagonal(nz_direction, True)
 
     # X edges
     nynz_direction = ny_direction + nz_direction
@@ -129,72 +129,81 @@ def triangulate_dual_contouring(dc_data_per_surface: DualContouringData):
     xy_direction = x_direction + y_direction
 
     # Stack all 12 directions
-    directions = np.dstack([nynz_direction, nyz_direction, ynz_direction, yz_direction,
-                            nxnz_direction, xnz_direction, nxz_direction, xz_direction,
-                            nxny_direction, nxy_direction, xny_direction, xy_direction])
+    directions = BackendTensor.tfnp.dstack([nynz_direction, nyz_direction, ynz_direction, yz_direction,
+                                            nxnz_direction, xnz_direction, nxz_direction, xz_direction,
+                                            nxny_direction, nxy_direction, xny_direction, xy_direction])
 
     # endregion
 
     valid_edg = valid_edges[valid_voxels][:, :]
-    valid_edg = BackendTensor.t.to_numpy(valid_edg)
-    
+    valid_edg = BackendTensor.tfnp.to_numpy(valid_edg)
+
     direction_each_edge = (directions * valid_edg)
 
     # Pick only edges with more than 2 voxels nearby
     three_neighbours = (directions * valid_edg).sum(axis=0) == 3
-    matrix_to_right_C_order = np.transpose((direction_each_edge * three_neighbours), (1, 2, 0))
-    indices = np.where(matrix_to_right_C_order)[2].reshape(-1, 3)
+    matrix_to_right_C_order = BackendTensor.tfnp.transpose((direction_each_edge * three_neighbours), (1, 2, 0))
+    indices = BackendTensor.tfnp.where(matrix_to_right_C_order)[2].reshape(-1, 3)
 
     indices_shift = indices
     indices_arrays.append(indices_shift)
-    indices_arrays_f = np.vstack(indices_arrays)
+    indices_arrays_f = BackendTensor.tfnp.vstack(indices_arrays)
 
     return indices_arrays_f
 
 
-def generate_dual_contouring_vertices(dc_data_per_stack: DualContouringData, slice_surface: slice, debug: bool = False) -> np.ndarray:
+def generate_dual_contouring_vertices(dc_data_per_stack: DualContouringData, slice_surface: slice, debug: bool = False):
     # @off
-    n_edges      = dc_data_per_stack.n_edges
-    valid_edges  = dc_data_per_stack.valid_edges
+    n_edges = dc_data_per_stack.n_edges
+    valid_edges = dc_data_per_stack.valid_edges
     valid_voxels = dc_data_per_stack.valid_voxels
-    xyz_on_edge  = dc_data_per_stack.xyz_on_edge[slice_surface]
-    gradients    = dc_data_per_stack.gradients[slice_surface]
+    xyz_on_edge = dc_data_per_stack.xyz_on_edge[slice_surface]
+    gradients = dc_data_per_stack.gradients[slice_surface]
     # @on
 
     # * Coordinates for all posible edges (12) and 3 dummy edges_normals in the center
-    edges_xyz = BackendTensor.t.zeros((n_edges, 15, 3), dtype=BackendTensor.dtype_obj)
+    edges_xyz = BackendTensor.tfnp.zeros((n_edges, 15, 3), dtype=BackendTensor.dtype_obj)
     valid_edges = valid_edges > 0
     edges_xyz[:, :12][valid_edges] = xyz_on_edge
 
     # Normals
-    edges_normals = BackendTensor.t.zeros((n_edges, 15, 3), dtype=BackendTensor.dtype_obj)
+    edges_normals = BackendTensor.tfnp.zeros((n_edges, 15, 3), dtype=BackendTensor.dtype_obj)
     edges_normals[:, :12][valid_edges] = gradients
 
-    if OLD_METHOD:=False:
+    if OLD_METHOD := False:
         # ! Moureze model does not seems to work with the new method
         # ! This branch is all nans at least with ch1_1 model
-        bias_xyz = np.copy(edges_xyz[:, :12])
-        isclose = np.isclose(bias_xyz, 0)
-        bias_xyz[isclose] = np.nan  # np zero values to nans
-        mass_points = np.nanmean(bias_xyz, axis=1)  # Mean ignoring nans
+        bias_xyz = BackendTensor.tfnp.copy(edges_xyz[:, :12])
+        isclose = BackendTensor.tfnp.isclose(bias_xyz, 0)
+        bias_xyz[isclose] = BackendTensor.tfnp.nan  # zero values to nans
+        mass_points = BackendTensor.tfnp.nanmean(bias_xyz, axis=1)  # Mean ignoring nans
     else:  # ? This is actually doing something
-        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)
-        mass_points = BackendTensor.t.array(mass_points)
+        bias_xyz = BackendTensor.tfnp.copy(edges_xyz[:, :12])
+        if BackendTensor.engine_backend == AvailableBackends.PYTORCH:
+            # PyTorch doesn't have masked arrays, so we'll use a different approach
+            mask = bias_xyz == 0
+            # Replace zeros with NaN for mean calculation
+            bias_xyz_masked = BackendTensor.tfnp.where(mask, float('nan'), bias_xyz)
+            mass_points = BackendTensor.tfnp.nanmean(bias_xyz_masked, axis=1)
+        else:
+            # NumPy approach with masked arrays
+            bias_xyz = BackendTensor.tfnp.to_numpy(bias_xyz)
+            import numpy as np
+            mask = bias_xyz == 0
+            masked_arr = np.ma.masked_array(bias_xyz, mask)
+            mass_points = masked_arr.mean(axis=1)
+            mass_points = BackendTensor.tfnp.array(mass_points)
 
     edges_xyz[:, 12] = mass_points
     edges_xyz[:, 13] = mass_points
     edges_xyz[:, 14] = mass_points
 
     BIAS_STRENGTH = 1
-    
-    bias_x = BackendTensor.t.array([BIAS_STRENGTH, 0, 0], dtype=BackendTensor.dtype_obj)
-    bias_y = BackendTensor.t.array([0, BIAS_STRENGTH, 0], dtype=BackendTensor.dtype_obj)
-    bias_z = BackendTensor.t.array([0, 0, BIAS_STRENGTH], dtype=BackendTensor.dtype_obj)
-    
+
+    bias_x = BackendTensor.tfnp.array([BIAS_STRENGTH, 0, 0], dtype=BackendTensor.dtype_obj)
+    bias_y = BackendTensor.tfnp.array([0, BIAS_STRENGTH, 0], dtype=BackendTensor.dtype_obj)
+    bias_z = BackendTensor.tfnp.array([0, 0, BIAS_STRENGTH], dtype=BackendTensor.dtype_obj)
+
     edges_normals[:, 12] = bias_x
     edges_normals[:, 13] = bias_y
     edges_normals[:, 14] = bias_z
@@ -208,14 +217,15 @@ def generate_dual_contouring_vertices(dc_data_per_stack: DualContouringData, sli
     b = (A * edges_xyz).sum(axis=2)
 
     if BackendTensor.engine_backend == AvailableBackends.PYTORCH:
-        transpose_shape = (2, 1)
+        transpose_shape = (2, 1, 0)  # For PyTorch: (batch, dim2, dim1)
     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, transpose_shape), b)
-    vertices = BackendTensor.t.einsum("ijk, ij->ik", term2, term3)
+        transpose_shape = (0, 2, 1)  # For NumPy: (batch, dim2, dim1)
+
+    import torch
+    term1 = BackendTensor.tfnp.einsum("ijk, ilj->ikl", A, BackendTensor.tfnp.transpose(A, transpose_shape))
+    term2 = BackendTensor.tfnp.linalg.inv(term1)
+    term3 = BackendTensor.tfnp.einsum("ijk,ik->ij", BackendTensor.tfnp.transpose(A, transpose_shape), b)
+    vertices = BackendTensor.tfnp.einsum("ijk, ij->ik", term2, term3)
 
     if debug:
         dc_data_per_stack.bias_center_mass = edges_xyz[:, 12:].reshape(-1, 3)
@@ -237,8 +247,8 @@ def evaluate(self, x):
         """Evaluates the function at a given point.
         This is what the solve method is trying to minimize.
         NB: Doesn't work with fixed axes."""
-        x = np.array(x)
-        return np.linalg.norm(np.matmul(self.A, x) - self.b)
+        x = BackendTensor.tfnp.array(x)
+        return BackendTensor.tfnp.linalg.norm(BackendTensor.tfnp.matmul(self.A, x) - self.b)
 
     def eval_with_pos(self, x):
         """Evaluates the QEF at a position, returning the same format solve does."""
@@ -248,15 +258,15 @@ def eval_with_pos(self, x):
     def make_3d(positions, normals):
         """Returns a QEF that measures the the error from a bunch of normals, each emanating
          from given positions"""
-        A = np.array(normals)
+        A = BackendTensor.tfnp.array(normals)
         b = [v[0] * n[0] + v[1] * n[1] + v[2] * n[2] for v, n in zip(positions, normals)]
         fixed_values = [None] * A.shape[1]
         return QEF(A, b, fixed_values)
 
     def solve(self):
         """Finds the point that minimizes the error of this QEF,
         and returns a tuple of the error squared and the point itself"""
-        result, residual, rank, s = np.linalg.lstsq(self.A, self.b)
+        result, residual, rank, s = BackendTensor.tfnp.linalg.lstsq(self.A, self.b)
         if len(residual) == 0:
             residual = self.evaluate(result)
         else: