[ENH] Add PyTorch backend support for dual contouring and octree operations (#21)

# [ENH] Adding PyTorch GPU Support

This PR adds support for PyTorch GPU acceleration by:

- Implementing a PyTorch backend for tensor operations throughout the codebase
- Adding GPU device detection and configuration in the backend tensor module
- Creating PyTorch-compatible implementations of key functions like `packbits`, `to_numpy`, and matrix operations
- Adapting octree generation and dual contouring algorithms to work with PyTorch tensors
- Ensuring proper tensor conversion between CPU and GPU when needed
- Implementing tensor-specific operations for both backends to maintain compatibility
- Adding proper memory management for contiguous arrays in PyTorch

The implementation ensures that operations can run on either CPU or GPU when using the PyTorch backend, with appropriate error handling when CUDA is not available but GPU is requested.

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/API/interp_single/_octree_generation.py

@@ -58,7 +58,7 @@ def interpolate_on_octree(interpolation_input: InterpolationInput, options: Inte
     return next_octree_level
 
 
-def _generate_corners(regular_grid: RegularGrid, level=1) -> np.ndarray:
+def _generate_corners_DEP(regular_grid: RegularGrid, level=1) -> np.ndarray:
     if regular_grid is None: raise ValueError("Regular grid is None")
 
     xyz_coord = regular_grid.values
@@ -79,3 +79,41 @@ def stack_left_right(a_edg, d_a):
 
     new_xyz = np.stack((x, y, z)).T
     return np.ascontiguousarray(new_xyz)
+
+def _generate_corners(regular_grid: RegularGrid, level=1):
+    if regular_grid is None:
+        raise ValueError("Regular grid is None")
+
+    # Convert to backend tensors
+    # xyz_coord = BackendTensor.tfnp.array(regular_grid.values)
+    # dxdydz = BackendTensor.tfnp.array(regular_grid.dxdydz)
+
+    xyz_coord = regular_grid.values
+    dxdydz = regular_grid.dxdydz
+
+    x_coord, y_coord, z_coord = xyz_coord[:, 0], xyz_coord[:, 1], xyz_coord[:, 2]
+    dx, dy, dz = dxdydz[0], dxdydz[1], dxdydz[2]
+
+    def stack_left_right(a_edg, d_a):
+        left = a_edg - d_a / level / 2
+        right = a_edg + d_a / level / 2
+        return BackendTensor.tfnp.stack([left, right], axis=1)
+
+    x_ = BackendTensor.tfnp.repeat(stack_left_right(x_coord, dx), 4, axis=1)
+    x = BackendTensor.tfnp.ravel(x_)
+
+    y_temp = BackendTensor.tfnp.repeat(stack_left_right(y_coord, dy), 2, axis=1)
+    y_ = BackendTensor.tfnp.tile(y_temp, (1, 2))
+    y = BackendTensor.tfnp.ravel(y_)
+
+    z_ = BackendTensor.tfnp.tile(stack_left_right(z_coord, dz), (1, 4))
+    z = BackendTensor.tfnp.ravel(z_)
+
+    new_xyz = BackendTensor.tfnp.stack([x, y, z], axis=1)
+
+    # For PyTorch, ensure contiguous memory (equivalent to np.ascontiguousarray)
+    if BackendTensor.engine_backend == AvailableBackends.PYTORCH:
+        if hasattr(new_xyz, 'contiguous'):
+            new_xyz = new_xyz.contiguous()
+
+    return new_xyz

gempy_engine/core/backend_tensor.py

@@ -107,9 +107,20 @@ def _change_backend(cls, engine_backend: AvailableBackends, use_pykeops: bool =
                 if (use_gpu):
                     cls.use_gpu = True
                     # cls.tensor_backend_pointer['active_backend'].set_default_device("cuda")
+                    # Check if CUDA is available
+                    if not pytorch_copy.cuda.is_available():
+                        raise RuntimeError("GPU requested but CUDA is not available in PyTorch")
+                    if False:
+                        # Set default device to CUDA
+                        cls.device = pytorch_copy.device("cuda")
+                        pytorch_copy.set_default_device("cuda")
+                        print(f"GPU enabled. Using device: {cls.device}")
+                        print(f"GPU device count: {pytorch_copy.cuda.device_count()}")
+                        print(f"Current GPU device: {pytorch_copy.cuda.current_device()}")
                 else:
                     cls.use_gpu = False
-
+                    cls.device = pytorch_copy.device("cpu")
+                    pytorch_copy.set_default_device("cpu")
             case (_):
                 raise AttributeError(
                     f"Engine Backend: {engine_backend} cannot be used because the correspondent library"
@@ -134,7 +145,7 @@ def describe_conf(cls):
 
     @classmethod
     def _wrap_pytorch_functions(cls):
-        from torch import sum, repeat_interleave
+        from torch import sum, repeat_interleave, isclose
         import torch
 
         def _sum(tensor, axis=None, dtype=None, keepdims=False):
@@ -155,6 +166,11 @@ def _array(array_like, dtype=None):
                 if dtype is None: return array_like
                 else: return array_like.type(dtype)
             else:
+                # Ensure numpy arrays are contiguous before converting to torch tensor
+                if isinstance(array_like, numpy.ndarray):
+                    if not array_like.flags.c_contiguous:
+                        array_like = numpy.ascontiguousarray(array_like)
+
                 return torch.tensor(array_like, dtype=dtype)
 
         def _concatenate(tensors, axis=0, dtype=None):
@@ -167,6 +183,95 @@ 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"):
+            """
+            Pack boolean values into uint8 bytes along the specified axis.
+            For a (4, n) tensor with axis=0, this packs every 4 bits into nibbles,
+            then pads to create full bytes.
+            """
+            # Convert to uint8 if boolean
+            if tensor.dtype == torch.bool:
+                tensor = tensor.to(torch.uint8)
+
+            if axis == 0:
+                # Pack along axis 0 (rows)
+                n_rows, n_cols = tensor.shape
+                n_output_rows = (n_rows + 7) // 8  # Round up to nearest byte boundary
+
+                # Pad with zeros if we don't have multiples of 8 rows
+                if n_rows % 8 != 0:
+                    padding_rows = 8 - (n_rows % 8)
+                    padding = torch.zeros(padding_rows, n_cols, dtype=torch.uint8, device=tensor.device)
+                    tensor = torch.cat([tensor, padding], dim=0)
+
+                # Reshape to group every 8 rows together: (n_output_rows, 8, n_cols)
+                tensor_reshaped = tensor.view(n_output_rows, 8, n_cols)
+
+                # Define bit positions (powers of 2)
+                if bitorder == "little":
+                    # Little endian: LSB first [1, 2, 4, 8, 16, 32, 64, 128]
+                    powers = torch.tensor([1, 2, 4, 8, 16, 32, 64, 128],
+                                          dtype=torch.uint8, device=tensor.device).view(1, 8, 1)
+                else:
+                    # Big endian: MSB first [128, 64, 32, 16, 8, 4, 2, 1]
+                    powers = torch.tensor([128, 64, 32, 16, 8, 4, 2, 1],
+                                          dtype=torch.uint8, device=tensor.device).view(1, 8, 1)
+
+                # Pack bits: multiply each bit by its power and sum along the 8-bit dimension
+                packed = (tensor_reshaped * powers).sum(dim=1)  # Shape: (n_output_rows, n_cols)
+
+                return packed
+
+            elif axis == 1:
+                # Pack along axis 1 (columns) 
+                n_rows, n_cols = tensor.shape
+                n_output_cols = (n_cols + 7) // 8
+
+                # Pad with zeros if needed
+                if n_cols % 8 != 0:
+                    padding_cols = 8 - (n_cols % 8)
+                    padding = torch.zeros(n_rows, padding_cols, dtype=torch.uint8, device=tensor.device)
+                    tensor = torch.cat([tensor, padding], dim=1)
+
+                # Reshape: (n_rows, n_output_cols, 8)
+                tensor_reshaped = tensor.view(n_rows, n_output_cols, 8)
+
+                # Define bit positions
+                if bitorder == "little":
+                    powers = torch.tensor([1, 2, 4, 8, 16, 32, 64, 128],
+                                          dtype=torch.uint8, device=tensor.device).view(1, 1, 8)
+                else:
+                    powers = torch.tensor([128, 64, 32, 16, 8, 4, 2, 1],
+                                          dtype=torch.uint8, device=tensor.device).view(1, 1, 8)
+
+                packed = (tensor_reshaped * powers).sum(dim=2)  # Shape: (n_rows, n_output_cols)
+                return packed
+
+            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
+
+        def _fill_diagonal(tensor, value):
+            """Fill the diagonal of a 2D tensor with the given value"""
+            if tensor.dim() != 2:
+                raise ValueError("fill_diagonal only supports 2D tensors")
+            diagonal_indices = torch.arange(min(tensor.size(0), tensor.size(1)))
+            tensor[diagonal_indices, diagonal_indices] = value
+            return tensor
 
         cls.tfnp.sum = _sum
         cls.tfnp.repeat = _repeat
@@ -175,7 +280,7 @@ def _transpose(tensor, axes=None):
         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)
@@ -185,6 +290,17 @@ def _transpose(tensor, axes=None):
         cls.tfnp.transpose = _transpose
         cls.tfnp.geomspace = lambda start, stop, step: torch.logspace(start, stop, step, base=10)
         cls.tfnp.abs = lambda tensor, dtype = None: tensor.abs().type(dtype) if dtype is not None else tensor.abs()
+        cls.tfnp.tile = lambda tensor, repeats: tensor.repeat(repeats)
+        cls.tfnp.ravel = lambda tensor: tensor.flatten()
+        cls.tfnp.packbits = _packbits
+        cls.tfnp.fill_diagonal = _fill_diagonal
+        cls.tfnp.isclose = lambda a, b, rtol=1e-05, atol=1e-08, equal_nan=False: isclose(
+            a,
+            torch.tensor(b, dtype=a.dtype, device=a.device),
+            rtol=rtol,
+            atol=atol,
+            equal_nan=equal_nan
+        )
 
     @classmethod
     def _wrap_pykeops_functions(cls):

gempy_engine/core/data/interp_output.py

@@ -127,8 +127,8 @@ def litho_faults_ids(self):
         faults_ids = BackendTensor.t.rint(self.faults_block)
 
         # Get the number of unique lithology IDs
-        multiplier = len(np.unique(litho_ids))
-
+        multiplier = len(BackendTensor.t.unique(litho_ids))
+    
         # Generate the unique IDs
         unique_ids = litho_ids + faults_ids * multiplier
         return unique_ids
@@ -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/core/data/regular_grid.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 
+from ..backend_tensor import BackendTensor
 from ..utils import _check_and_convert_list_to_array, cast_type_inplace
 from .kernel_classes.server.input_parser import GridSchema
 
@@ -27,7 +28,7 @@ def __post_init__(self):
 
         self._create_regular_grid_3d()
 
-        self.original_values = self.values.copy()
+        self.original_values = BackendTensor.t.copy(self.values)
 
     @property
     def dx(self):
@@ -196,5 +197,6 @@ def _create_regular_grid_3d(self):
         g = np.meshgrid(*coords, indexing="ij")
         values = np.vstack(tuple(map(np.ravel, g))).T.astype("float64")
         values = np.ascontiguousarray(values)
+        values = BackendTensor.tfnp.array(values, dtype=BackendTensor.dtype_obj)
 
         self.values = values

gempy_engine/core/utils.py

@@ -12,7 +12,11 @@ def cast_type_inplace(data_instance: Any, requires_grad:bool = False):
             case (gempy_engine.config.AvailableBackends.numpy):
                 data_instance.__dict__[key] = val.astype(BackendTensor.dtype)
             case (gempy_engine.config.AvailableBackends.PYTORCH):
-                tensor = BackendTensor.t.from_numpy(val.astype(BackendTensor.dtype))
+                # tensor = BackendTensor.t.from_numpy(val.astype(BackendTensor.dtype))
+                # if (BackendTensor.use_gpu):
+                #     tensor = tensor.cuda()
+
+                tensor = BackendTensor.tfnp.array(val, dtype=BackendTensor.dtype_obj)
                 tensor.requires_grad = requires_grad
                 data_instance.__dict__[key] = tensor