[WIP] Trying to do the triangulation in Torch

Author: Leguark

gempy_engine/API/dual_contouring/_dual_contouring.py

@@ -54,10 +54,30 @@ def compute_dual_contouring(dc_data_per_stack: DualContouringData, left_right_co
             edges_normals[:] = np.nan
             edges_normals[valid_edges] = dc_data_per_stack.gradients[slice_object]
 
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore", category=RuntimeWarning)
-                voxel_normal  = np.nanmean(edges_normals, axis=1)
-                voxel_normal  = voxel_normal[(~np.isnan(voxel_normal).any(axis=1))]  # drop nans
+            if LEGACY:=False:
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore", category=RuntimeWarning)
+                    voxel_normal  = np.nanmean(edges_normals, axis=1)
+                    voxel_normal  = voxel_normal[(~np.isnan(voxel_normal).any(axis=1))]  # drop nans
+                    pass 
+            else:
+                # Assuming edges_normals is a PyTorch tensor
+                nan_mask = BackendTensor.t.isnan(edges_normals)
+                valid_count = (~nan_mask).sum(dim=1)
+
+                # Replace NaNs with 0 for sum calculation
+                safe_normals = edges_normals.clone()
+                safe_normals[nan_mask] = 0
+
+                # Compute the sum of non-NaN elements
+                sum_normals = BackendTensor.t.sum(safe_normals, 1)
+
+                # Calculate the mean, avoiding division by zero
+                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]
+                
 
             valid_voxels = dc_data_per_surface.valid_voxels
             indices = triangulate(

gempy_engine/API/interp_single/_interp_single_feature.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from gempy_engine.config import AvailableBackends
+from gempy_engine.config import AvailableBackends, COMPUTE_GRADIENTS
 from ...core.backend_tensor import BackendTensor
 from ._interp_scalar_field import interpolate_scalar_field
 from ...core.data import SurfacePoints, SurfacePointsInternals, Orientations, OrientationsInternals, TensorsStructure
@@ -26,7 +26,10 @@ def interpolate_feature(interpolation_input: InterpolationInput,
                         external_segment_funct: Optional[Callable[[np.ndarray], float]] = None,
                         clean_buffer: bool = True) -> ScalarFieldOutput:
 
-    grid = copy.deepcopy(interpolation_input.grid)
+    if BackendTensor.engine_backend is not AvailableBackends.PYTORCH and COMPUTE_GRADIENTS is False:
+        grid = copy.deepcopy(interpolation_input.grid)
+    else:
+        grid = interpolation_input.grid
 
     # region Interpolate scalar field
     xyz = solver_input.xyz_to_interpolate

gempy_engine/API/interp_single/_octree_generation.py

@@ -21,6 +21,8 @@ def interpolate_on_octree(interpolation_input: InterpolationInput, options: Inte
 
     if BackendTensor.engine_backend is not AvailableBackends.PYTORCH and COMPUTE_GRADIENTS is False:
         temp_interpolation_input = copy.deepcopy(interpolation_input)
+    else:
+        temp_interpolation_input = interpolation_input
 
     # * Interpolate - centers
     output_0_centers: List[InterpOutput] = interpolate_all_fields(temp_interpolation_input, options, data_shape)  # interpolate - centers

gempy_engine/API/interp_single/interp_features.py

@@ -2,6 +2,8 @@
 from typing import List
 
 import gempy_engine.core.data.tensors_structure
+from gempy_engine.config import AvailableBackends, COMPUTE_GRADIENTS
+from ...core.backend_tensor import BackendTensor
 from ...core.data.grid import Grid
 from ...core import data
 from ...core.data import InterpolationOptions
@@ -45,8 +47,12 @@ def interpolate_n_octree_levels(interpolation_input: InterpolationInput, options
 
 def interpolate_all_fields_no_octree(interpolation_input: InterpolationInput, options: InterpolationOptions,
                                      data_descriptor: InputDataDescriptor) -> List[InterpOutput]:
-    interpolation_input = copy.deepcopy(interpolation_input)
-    return ms.interpolate_all_fields(interpolation_input, options, data_descriptor)
+    if BackendTensor.engine_backend is not AvailableBackends.PYTORCH and COMPUTE_GRADIENTS is False:
+        temp_interpolation_input = copy.deepcopy(interpolation_input)
+    else:
+        temp_interpolation_input = interpolation_input
+
+    return ms.interpolate_all_fields(temp_interpolation_input, options, data_descriptor)
 
 
 # region testing

gempy_engine/core/backend_tensor.py

@@ -169,7 +169,19 @@ def _array(array_like, dtype=None):
                 dtype = getattr(torch, dtype)
 
             return torch.tensor(array_like, dtype=dtype)
-
+        
+        def _concatenate(tensors, axis=0, dtype=None):
+            # Switch if tensor is numpy array or a torch tensor
+            match type(tensors[0]):
+                case numpy.ndarray:
+                    return numpy.concatenate(tensors, axis=axis)
+                case torch.Tensor:
+                    return torch.cat(tensors, dim=axis)
+        
+        def _transpose(tensor, axes=None):
+            return tensor.transpose(axes[0], axes[1])
+        
+        
         cls.tfnp.sum = _sum
         cls.tfnp.repeat = _repeat
         cls.tfnp.expand_dims = lambda tensor, axis: tensor
@@ -182,7 +194,8 @@ def _array(array_like, dtype=None):
         cls.tfnp.rint = lambda tensor: tensor.round().type(torch.int32)
         cls.tfnp.vstack = lambda tensors: torch.cat(tensors, dim=0)
         cls.tfnp.copy = lambda tensor: tensor.clone()
-    
+        cls.tfnp.concatenate = _concatenate
+        cls.tfnp.transpose = _transpose
 
 
     @classmethod

gempy_engine/core/data/grid.py

@@ -46,7 +46,7 @@ def values(self) -> np.ndarray:
         if self.geophysics_grid is not None:
             values.append(self.geophysics_grid.values)
 
-        values_array = np.concatenate(values, dtype=BackendTensor.dtype)
+        values_array = BackendTensor.t.concatenate(values, dtype=BackendTensor.dtype)
         values_array = BackendTensor.t.array(values_array)
 
         return values_array

gempy_engine/modules/dual_contouring/dual_contouring_interface.py

@@ -173,7 +173,7 @@ 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 = np.copy(edges_xyz[:, :12])
+        bias_xyz = BackendTensor.t.copy(edges_xyz[:, :12]).detach().numpy()
         mask = bias_xyz == 0
         masked_arr = np.ma.masked_array(bias_xyz, mask)
         mass_points = masked_arr.mean(axis=1)
@@ -200,10 +200,11 @@ 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 = np.einsum("ijk, ilj->ikl", A, np.transpose(A, (0, 2, 1)))
-    term2 = np.linalg.inv(term1)
-    term3 = np.einsum("ijk,ik->ij", np.transpose(A, (0, 2, 1)), b)
-    vertices = np.einsum("ijk, ij->ik", term2, term3)
+    
+    term1 = BackendTensor.t.einsum("ijk, ilj->ikl", A, BackendTensor.t.transpose(A, (2, 1)))
+    term2 = BackendTensor.t.linalg.inv(term1)
+    term3 = BackendTensor.t.einsum("ijk,ik->ij", BackendTensor.t.transpose(A, (2, 1)), b)
+    vertices = BackendTensor.t.einsum("ijk, ij->ik", term2, term3)
 
     if debug:
         dc_data_per_stack.bias_center_mass = edges_xyz[:, 12:].reshape(-1, 3)