diff --git a/gempy_engine/API/interp_single/_interp_scalar_field.py b/gempy_engine/API/interp_single/_interp_scalar_field.py index be749af..34ff698 100644 --- a/gempy_engine/API/interp_single/_interp_scalar_field.py +++ b/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 diff --git a/gempy_engine/core/backend_tensor.py b/gempy_engine/core/backend_tensor.py index d734e9b..d8f5536 100644 --- a/gempy_engine/core/backend_tensor.py +++ b/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") diff --git a/gempy_engine/modules/evaluator/symbolic_evaluator.py b/gempy_engine/modules/evaluator/symbolic_evaluator.py index 6fa3399..9712c00 100644 --- a/gempy_engine/modules/evaluator/symbolic_evaluator.py +++ b/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: diff --git a/tests/benchmark/profile_runner.py b/tests/benchmark/profile_runner.py index 6ab2e39..b0f6cd5 100644 --- a/tests/benchmark/profile_runner.py +++ b/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( diff --git a/tests/conftest.py b/tests/conftest.py index 180b821..4bbb667 100644 --- a/tests/conftest.py +++ b/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: diff --git a/tests/test_common/test_modules/test_kernel_constructor/test_kernel_constructor.py b/tests/test_common/test_modules/test_kernel_constructor/test_kernel_constructor.py index d7f357c..2c05b45 100644 --- a/tests/test_common/test_modules/test_kernel_constructor/test_kernel_constructor.py +++ b/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) diff --git a/tests/test_common/test_modules/test_kernel_constructor/test_kernel_constructor_2.py b/tests/test_common/test_modules/test_kernel_constructor/test_kernel_constructor_2.py index a0e6068..0bdfe83 100644 --- a/tests/test_common/test_modules/test_kernel_constructor/test_kernel_constructor_2.py +++ b/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]