[BUG/CLN] Refactor tensor handling and cleanup lazy weight logic

gempy_engine/API/interp_single/_interp_scalar_field.py

@@ -43,11 +43,11 @@ def interpolate_scalar_field(solver_input: SolverInput, options: InterpolationOp
             weights_cached = None
         case _:
             raise ValueError("Cache mode not recognized")
-
 
+
+    BackendTensor.pykeops_enabled = False
     match weights_cached:
         case None:
-            foo = solver_input.weights_x0
             weights = _solve_and_store_weights(
                 solver_input=solver_input,
                 kernel_options=options.kernel_options,
@@ -70,6 +70,7 @@ def interpolate_scalar_field(solver_input: SolverInput, options: InterpolationOp
 
     # endregion
 
+    BackendTensor.pykeops_enabled = BackendTensor.use_pykeops
     exported_fields: ExportedFields = _evaluate_sys_eq(solver_input, weights, options)
 
     return weights, exported_fields

gempy_engine/core/backend_tensor.py

@@ -22,6 +22,7 @@ class BackendTensor:
     engine_backend: AvailableBackends
 
     pykeops_enabled: bool
+    use_pykeops: bool = False
     use_gpu: bool = True
     dtype: str = DEFAULT_TENSOR_DTYPE
     dtype_obj: Union[str, "torch.dtype"] = DEFAULT_TENSOR_DTYPE
@@ -46,10 +47,10 @@ def get_backend_string(cls) -> str:
 
     @classmethod
     def change_backend_gempy(cls, engine_backend: AvailableBackends, use_gpu: bool = True, dtype: Optional[str] = None):
-        cls._change_backend(engine_backend, pykeops_enabled=PYKEOPS, use_gpu=use_gpu, dtype=dtype)
+        cls._change_backend(engine_backend, use_pykeops=PYKEOPS, use_gpu=use_gpu, dtype=dtype)
 
     @classmethod
-    def _change_backend(cls, engine_backend: AvailableBackends, pykeops_enabled: bool = False, use_gpu: bool = True, dtype: Optional[str] = None):
+    def _change_backend(cls, engine_backend: AvailableBackends, use_pykeops: bool = False, use_gpu: bool = True, dtype: Optional[str] = None):
         cls.dtype = DEFAULT_TENSOR_DTYPE if dtype is None else dtype
         cls.dtype_obj = cls.dtype
         match engine_backend:
@@ -71,20 +72,20 @@ def _change_backend(cls, engine_backend: AvailableBackends, pykeops_enabled: boo
 
                 cls._wrap_numpy_functions()
 
-                match (pykeops_enabled, is_pykeops_installed, use_gpu):
+                match (use_pykeops, is_pykeops_installed, use_gpu):
                     case (True, True, True):
-                        cls.pykeops_enabled = True
+                        cls.use_pykeops = True
                         cls.use_gpu = True
                         cls._wrap_pykeops_functions()
                     case (True, True, False):
-                        cls.pykeops_enabled = True
+                        cls.use_pykeops = True
                         cls.use_gpu = False
                         cls._wrap_pykeops_functions()
                     case (True, False, _):
                         raise AttributeError(
                             f"Engine Backend: {engine_backend} cannot be used because the correspondent library is not installed: pykeops")
                     case (False, _, _):
-                        cls.pykeops_enabled = False
+                        cls.use_pykeops = False
                         cls.use_gpu = False
 
             case (engine_backend.PYTORCH):
@@ -98,8 +99,8 @@ def _change_backend(cls, engine_backend: AvailableBackends, pykeops_enabled: boo
                 cls.dtype_obj = pytorch_copy.float32 if cls.dtype == "float32" else pytorch_copy.float64
                 cls.tensor_types = pytorch_copy.Tensor
 
-                cls.pykeops_enabled = pykeops_enabled  # TODO: Make this compatible with pykeops
-                if (pykeops_enabled):
+                cls.use_pykeops = use_pykeops  # TODO: Make this compatible with pykeops
+                if (use_pykeops):
                     import pykeops
                     cls._wrap_pykeops_functions()
 
@@ -209,6 +210,8 @@ def _sum(tensor, axis=None, dtype=None, keepdims=False):
                 case pykeops.numpy.LazyTensor() | pykeops.torch.LazyTensor():
                     return tensor.sum(axis)
                 case torch.Tensor() if torch_available:
+                    if isinstance(dtype, str):
+                        dtype = getattr(torch, dtype)
                     return tensor.sum(axis, keepdims=keepdims, dtype=dtype)
                 case _:
                     raise TypeError("Unsupported tensor type")

gempy_engine/modules/evaluator/symbolic_evaluator.py

@@ -34,12 +34,12 @@ def symbolic_evaluator(solver_input: SolverInput, weights: np.ndarray, options:
         eval_gx_kernel = yield_evaluation_grad_kernel(solver_input, options.kernel_options, axis=0)
         eval_gy_kernel = yield_evaluation_grad_kernel(solver_input, options.kernel_options, axis=1)
 
-        gx_field = (eval_gx_kernel.T * LazyTensor(weights, axis=1)).sum(axis=1, backend=backend_string).reshape(-1)
-        gy_field = (eval_gy_kernel.T * LazyTensor(weights, axis=1)).sum(axis=1, backend=backend_string).reshape(-1)
+        gx_field = (eval_gx_kernel.T * lazy_weights).sum(axis=1, backend=backend_string).reshape(-1)
+        gy_field = (eval_gy_kernel.T * lazy_weights).sum(axis=1, backend=backend_string).reshape(-1)
 
         if options.number_dimensions == 3:
             eval_gz_kernel = yield_evaluation_grad_kernel(solver_input, options.kernel_options, axis=2)
-            gz_field = (eval_gz_kernel.T * LazyTensor(weights, axis=1)).sum(axis=1, backend=backend_string).reshape(-1)
+            gz_field = (eval_gz_kernel.T * lazy_weights).sum(axis=1, backend=backend_string).reshape(-1)
         elif options.number_dimensions == 2:
             gz_field = None
         else:

tests/benchmark/profile_runner.py

@@ -24,7 +24,7 @@ def profile_moureze_model():
     BackendTensor._change_backend(
         engine_backend=AvailableBackends.numpy,
         use_gpu=False,
-        pykeops_enabled=False
+        use_pykeops=False
     )
 
     model = moureze_model_factory(

tests/conftest.py

@@ -24,7 +24,7 @@
 BackendTensor._change_backend(
     engine_backend=backend,
     use_gpu=use_gpu,
-    pykeops_enabled=pykeops_enabled
+    use_pykeops=pykeops_enabled
 )
 
 try:

tests/test_common/test_modules/test_kernel_constructor/test_kernel_constructor.py

@@ -185,7 +185,7 @@ def _compare_covariance_item_numpy_pykeops(self, preprocess_data, item, cov_func
         sp_internals, ori_internals, options = preprocess_data
 
         # numpy
-        BackendTensor._change_backend(AvailableBackends.numpy, pykeops_enabled=False)
+        BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=False)
         solver_input = SolverInput(sp_internals, ori_internals)
         kernel_data = cov_vectors_preparation(solver_input, options.kernel_options)
         c_n = cov_func(kernel_data, options, item=item)
@@ -198,7 +198,7 @@ def _compare_covariance_item_numpy_pykeops(self, preprocess_data, item, cov_func
         c_n_sum = c_n.sum(0).reshape(-1, 1)
 
         # pykeops
-        BackendTensor._change_backend(AvailableBackends.numpy, pykeops_enabled=True)
+        BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=True)
         kernel_data = cov_vectors_preparation(solver_input, options.kernel_options)
         c_k = cov_func(kernel_data, options, item=item)
         c_k_sum = c_n.sum(0).reshape(-1, 1)

tests/test_common/test_modules/test_kernel_constructor/test_kernel_constructor_2.py

@@ -61,7 +61,7 @@
 class TestCompareWithGempy_v2:
     @pytest.fixture(scope="class")
     def internals(self, simple_model):
-        BackendTensor._change_backend(AvailableBackends.numpy, pykeops_enabled=False)
+        BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=False)
 
         surface_points = simple_model[0]
         orientations = simple_model[1]