diff --git a/gempy_engine/API/model/model_api.py b/gempy_engine/API/model/model_api.py index 0404d5e8..ffc8b83e 100644 --- a/gempy_engine/API/model/model_api.py +++ b/gempy_engine/API/model/model_api.py @@ -18,6 +18,7 @@ from ...core.data.solutions import Solutions from ...core.utils import gempy_profiler_decorator from ...core.exceptions import GemPyEngineInputError +from ...core.data.options.temp_interpolation_values import TempInterpolationValues from ...modules.geophysics.fw_gravity import compute_gravity from ...modules.geophysics.fw_magnetic import compute_tmi from ...modules.weights_cache.weights_cache_interface import WeightCache @@ -26,6 +27,10 @@ @gempy_profiler_decorator def compute_model(interpolation_input: InterpolationInput, options: InterpolationOptions, data_descriptor: InputDataDescriptor, *, geophysics_input: Optional[GeophysicsInput] = None) -> Solutions: + # Octree progress and cache timestamps belong to this computation. Keeping + # them on a shared options object lets concurrent requests change each + # other's active octree level. + options = options.model_copy(update={"temp_interpolation_values": TempInterpolationValues()}) try: WeightCache.initialize_cache_dir() options.temp_interpolation_values.start_computation_ts = int(time.time()) diff --git a/gempy_engine/API/server/__init__.py b/gempy_engine/API/server/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/gempy_engine/API/server/_process_output.py b/gempy_engine/API/server/_process_output.py deleted file mode 100644 index 431a7aac..00000000 --- a/gempy_engine/API/server/_process_output.py +++ /dev/null @@ -1,130 +0,0 @@ -import json -import logging -from typing import List - -import numpy as np -import pandas as pd -import subsurface -from subsurface import UnstructuredData - -from gempy_engine.core.data.dual_contouring_mesh import DualContouringMesh -from gempy_engine.core.data.solutions import Solutions - -PLOT_SUBSURFACE_OBJECT = False - - -def process_output( - meshes: List[DualContouringMesh], - n_stack: int, - solutions: Solutions, - logger: logging.Logger -) -> bytes: - """ - Process model outputs into binary format for transmission. - - Args: - meshes: List of dual contouring meshes - n_stack: Number of stacks in the model - solutions: Model solutions containing octree data - logger: Logger instance - - Returns: - Binary data containing serialized meshes and octrees - """ - # Serialize meshes - unstructured_data_meshes: UnstructuredData = _meshes_to_unstruct(meshes, logger) - if PLOT_SUBSURFACE_OBJECT: - _plot_subsurface_object(unstructured_data_meshes) - body_meshes = unstructured_data_meshes.to_binary() - - # Serialize octrees - unstructured_data_volume = subsurface.UnstructuredData.from_array( - vertex=solutions.octrees_output[0].grid.values, - cells="points", - cells_attr=pd.DataFrame( - data=solutions.octrees_output[0].last_output_center.ids_block, - columns=['id'] - ) - ) - body_volume = unstructured_data_volume.to_binary() - - # Serialize global header and combine data - body = body_meshes + body_volume - global_header = json.dumps({"mesh_size": len(body_meshes), "octree_size": len(body_volume)}) - global_header_bytes = global_header.encode('utf-8') - global_header_length = len(global_header_bytes) - global_header_length_bytes = global_header_length.to_bytes(4, byteorder='little') - - return global_header_length_bytes + global_header_bytes + body - - -def _meshes_to_unstruct( - meshes: List[DualContouringMesh], - logger: logging.Logger -) -> UnstructuredData: - """ - Convert a list of dual contouring meshes to an unstructured data format. - - Args: - meshes: List of dual contouring meshes - logger: Logger instance - - Returns: - Unstructured data representation of the meshes - """ - n_meshes = len(meshes) - logger.debug(f"Number of meshes: {n_meshes}") - logger.debug(f"Mesh1TriShape: {meshes[0].edges.shape}") - - # Prepare the vertex array - vertex_array = np.concatenate([meshes[i].vertices for i in range(n_meshes)]) - - # Check for edge uniqueness (debugging) - simplex_array = np.concatenate([meshes[i].edges for i in range(n_meshes)]) - unc, count = np.unique(simplex_array, axis=0, return_counts=True) - logger.debug(f"edges shape {simplex_array.shape}") - - # Prepare the simplex array with proper indexing - simplex_array = meshes[0].edges - for i in range(1, n_meshes): - adder = np.max(meshes[i - 1].edges) + 1 - logger.debug(f"triangle counts adder: {adder}") - add_mesh = meshes[i].edges + adder - simplex_array = np.append(simplex_array, add_mesh, axis=0) - - logger.debug(f"edges shape {simplex_array.shape}") - - # Prepare the cells_attr array - ids_array = np.ones(simplex_array.shape[0]) - l0 = 0 - id_value = 1 - for mesh in meshes: - l1 = l0 + mesh.edges.shape[0] - logger.debug(f"l0 {l0} l1 {l1} id {id_value}") - ids_array[l0:l1] = id_value - l0 = l1 - id_value += 1 - logger.debug(f"ids_array count: {np.unique(ids_array)}") - - # Create the unstructured data - unstructured_data = subsurface.UnstructuredData.from_array( - vertex=vertex_array, - cells=simplex_array, - cells_attr=pd.DataFrame(ids_array, columns=['id']) - ) - - return unstructured_data - - -def _plot_subsurface_object(unstructured_data: UnstructuredData) -> None: - """ - Visualize the unstructured data using subsurface. - - Args: - unstructured_data: The unstructured data to visualize - """ - print(unstructured_data) - obj = subsurface.TriSurf(unstructured_data) - pv_unstruct = subsurface.visualization.to_pyvista_mesh(obj) - print(pv_unstruct) - subsurface.visualization.pv_plot([pv_unstruct]) \ No newline at end of file diff --git a/gempy_engine/API/server/_server_functions.py b/gempy_engine/API/server/_server_functions.py deleted file mode 100644 index aa6524fc..00000000 --- a/gempy_engine/API/server/_server_functions.py +++ /dev/null @@ -1,65 +0,0 @@ -import logging -import sys -from typing import Tuple - -from gempy_engine.core.data.input_data_descriptor import InputDataDescriptor -from gempy_engine.core.data.interpolation_input import InterpolationInput -from gempy_engine.core.data.kernel_classes.server.input_parser import GemPyInput - - -def setup_logger() -> logging.Logger: - """ - Configure and set up a logger for the application. - - Returns: - logging.Logger: Configured logger instance with console and file handlers - """ - logger = logging.getLogger("my-fastapi-app") - logger.setLevel(logging.DEBUG) - - formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") - - # Set up console handler for logging - ch = logging.StreamHandler(sys.stdout) - ch.setLevel(logging.DEBUG) - ch.setFormatter(formatter) - - # Set up file handler for logging - fh = logging.FileHandler("app.log") - fh.setLevel(logging.DEBUG) - fh.setFormatter(formatter) - - # Add the handlers to the logger - logger.addHandler(ch) - logger.addHandler(fh) - - return logger - - -def process_input( - gempy_input: GemPyInput, - logger: logging.Logger -) -> Tuple[InputDataDescriptor, InterpolationInput, int]: - """ - Process the GemPy input data to prepare it for model computation. - - Args: - gempy_input: Input data for the GemPy model - logger: Logger instance for recording processing information - - Returns: - Tuple containing: - - InputDataDescriptor: Structure descriptor for the model - - InterpolationInput: Prepared interpolation input data - - int: Number of stacks in the model - """ - logger.debug(f"Input grid: {gempy_input.interpolation_input.grid}") - - interpolation_input: InterpolationInput = InterpolationInput.from_schema(gempy_input.interpolation_input) - input_data_descriptor: InputDataDescriptor = InputDataDescriptor.from_schema(gempy_input.input_data_descriptor) - n_stack = len(input_data_descriptor.stack_structure.masking_descriptor) - - logger.debug(f"masking descriptor: {input_data_descriptor.stack_structure.masking_descriptor}") - logger.debug(f"stack structure: {input_data_descriptor.stack_structure}") - - return input_data_descriptor, interpolation_input, n_stack \ No newline at end of file diff --git a/gempy_engine/API/server/main_server_pro.py b/gempy_engine/API/server/main_server_pro.py deleted file mode 100644 index daaf940b..00000000 --- a/gempy_engine/API/server/main_server_pro.py +++ /dev/null @@ -1,110 +0,0 @@ -import fastapi -from fastapi import FastAPI -from fastapi.responses import Response - -from gempy_engine.API.model.model_api import compute_model -from gempy_engine.core.data import InterpolationOptions -from gempy_engine.core.data.input_data_descriptor import InputDataDescriptor -from gempy_engine.core.data.interpolation_input import InterpolationInput -from gempy_engine.core.data.kernel_classes.kernel_functions import AvailableKernelFunctions -from gempy_engine.core.data.kernel_classes.server.input_parser import GemPyInput -from gempy_engine.core.data.solutions import Solutions -from ._process_output import process_output -from ._server_functions import process_input, setup_logger - -# Optional visualization dependencies -try: - import pyvista as pv - from test.helper_functions_pyvista import plot_octree_pyvista, plot_dc_meshes, plot_pyvista - - VISUALIZATION_AVAILABLE = True -except ImportError: - VISUALIZATION_AVAILABLE = False - -# Create FastAPI application -gempy_engine_App = FastAPI(debug=True) -logger = setup_logger() - -# Default interpolation options -range_ = 1 -default_interpolation_options: InterpolationOptions = InterpolationOptions.from_args( - range=range_, - c_o=(range_ ** 2) / 14 / 3, - number_octree_levels=4, - kernel_function=AvailableKernelFunctions.cubic, - mesh_extraction=True -) - - -@gempy_engine_App.post("/") -def compute_gempy_model(gempy_input: GemPyInput) -> Response: - """ - Process GemPy input model and return serialized results. - - Args: - gempy_input: Input data for the GemPy model - - Returns: - Binary response containing serialized model output - """ - logger.info("Running GemPy Engine") - - - interpolation_input: InterpolationInput - input_data_descriptor: InputDataDescriptor - # Process input data - input_data_descriptor, interpolation_input, n_stack = process_input( - gempy_input=gempy_input, - logger=logger - ) - - # Compute model - solutions = _compute_model( - interpolation_input=interpolation_input, - options=default_interpolation_options, - structure=input_data_descriptor - ) - logger.info("Finished computing model") - - # Process output - body = process_output( - meshes=solutions.dc_meshes, - n_stack=n_stack, - solutions=solutions, - logger=logger - ) - - logger.info("Finished running GemPy Engine") - return fastapi.Response(content=body, media_type='application/octet-stream') - - -def _compute_model( - interpolation_input: InterpolationInput, - options: InterpolationOptions, - structure: InputDataDescriptor -) -> Solutions: - """ - Compute the GemPy model using the provided inputs. - - Args: - interpolation_input: Input data for interpolation - options: Interpolation options - structure: Data descriptor for model structure - - Returns: - Computed solutions - """ - solutions = compute_model(interpolation_input, options, structure) - - # Optional visualization (disabled by default) - if VISUALIZATION_AVAILABLE and False: - from test import helper_functions_pyvista - helper_functions_pyvista.plot_pyvista( - octree_list=None, - dc_meshes=solutions.dc_meshes, - gradients=interpolation_input.orientations.dip_gradients, - gradient_pos=interpolation_input.orientations.dip_gradients, - v_just_points=interpolation_input.surface_points.sp_coords - ) - - return solutions \ No newline at end of file diff --git a/gempy_engine/core/backend_tensor.py b/gempy_engine/core/backend_tensor.py index 941e0e3e..f3f04c69 100644 --- a/gempy_engine/core/backend_tensor.py +++ b/gempy_engine/core/backend_tensor.py @@ -185,6 +185,14 @@ def describe_conf(cls): print(f"\n Using gpu: {cls.use_gpu}. \n") print(f"\n Using pykeops: {cls.pykeops_enabled}. \n") + @classmethod + def arange(cls, stop, *, dtype=None): + if cls.engine_backend == AvailableBackends.PYTORCH: + if isinstance(dtype, str): + dtype = getattr(torch, dtype) + return torch.arange(stop, dtype=dtype, device=cls.device) + return cls.t.arange(stop, dtype=dtype) + @classmethod def _wrap_pytorch_functions(cls): import torch @@ -205,8 +213,12 @@ def _sum(tensor, axis=None, dtype=None, keepdims=False): def _repeat(tensor, n_repeats, axis=None): if not isinstance(tensor, torch.Tensor): tensor = torch.as_tensor(tensor, device=cls.device) + elif tensor.device != cls.device: + tensor = tensor.to(cls.device) if not isinstance(n_repeats, torch.Tensor): n_repeats = torch.as_tensor(n_repeats, device=cls.device) + elif n_repeats.device != cls.device: + n_repeats = n_repeats.to(cls.device) return _true_torch_repeat_interleave(tensor, n_repeats, dim=axis) def _array(array_like, dtype=None): diff --git a/gempy_engine/modules/data_preprocess/_input_preparation.py b/gempy_engine/modules/data_preprocess/_input_preparation.py index aa82042c..12fbdc23 100644 --- a/gempy_engine/modules/data_preprocess/_input_preparation.py +++ b/gempy_engine/modules/data_preprocess/_input_preparation.py @@ -38,7 +38,7 @@ def surface_points_preprocess(sp_input: SurfacePoints, tensors_structure: Tensor ref_points_repeated = b.t.repeat(ref_points, number_repetitions, 0) # ref_points shape: (1, 3) ref_nugget_repeated = b.t.repeat(ref_nugget, number_repetitions, 0) surface_ids = b.t.repeat( - b.t.arange(tensors_structure.n_surfaces, dtype=rest_nugget.dtype), + b.arange(tensors_structure.n_surfaces, dtype=rest_nugget.dtype), number_repetitions, 0, ) diff --git a/gempy_engine/modules/evaluator/symbolic_evaluator.py b/gempy_engine/modules/evaluator/symbolic_evaluator.py index a29a19e7..b0af7557 100644 --- a/gempy_engine/modules/evaluator/symbolic_evaluator.py +++ b/gempy_engine/modules/evaluator/symbolic_evaluator.py @@ -146,6 +146,21 @@ def _build_block_sparse_ranges(M_sizes: list[int], N_sizes: list[int]): return numpy_ranges +def _validate_stacked_dimensions(eval_kernel, weights, M_sizes: list[int], N_sizes: list[int]) -> None: + """Fail before PyKeOps when block ranges do not describe the lazy kernel.""" + expected_i = sum(N_sizes) + expected_j = sum(M_sizes) + kernel_shape = tuple(eval_kernel.shape) + weights_size = weights.shape[0] + + if kernel_shape[0] != expected_i or kernel_shape[1] != expected_j or weights_size != expected_i: + raise ValueError( + "Inconsistent stacked PyKeOps dimensions: " + f"kernel={kernel_shape}, weights={weights_size}, " + f"range dimensions=({expected_i}, {expected_j})." + ) + + def symbolic_evaluator_optimized_stacked( eval_inputs: list[EvaluatorInput], weights_list: list[np.ndarray], @@ -236,6 +251,7 @@ def _run_prep(args): # Concatenate weights all_weights = BackendTensor.t.concatenate(weights_list, axis=0) all_weights = BackendTensor.t.tile(all_weights, tile_factor) + _validate_stacked_dimensions(eval_kernel, all_weights, M_sizes, N_sizes) if BackendTensor.engine_backend == gempy_engine.config.AvailableBackends.numpy: from pykeops.numpy import LazyTensor @@ -276,13 +292,11 @@ def _run_prep(args): except TypeError: raise ValueError("Failed to compute symbolic evaluation with PyKeOps. Ensure that all_weights and eval_kernel are compatible for lazy tensor operations.") - # For torch - # all_results_concat = all_results_concat.to("cpu") - all_results_split = BackendTensor.t.split(all_results_concat, M_sizes) - - # For numpy - # split_indices = np.cumsum(M_sizes)[:-1] - # all_results_split = np.split(all_results_concat, split_indices) + if BackendTensor.engine_backend == gempy_engine.config.AvailableBackends.numpy: + split_indices = np.cumsum(M_sizes)[:-1] + all_results_split = np.split(all_results_concat, split_indices) + else: + all_results_split = BackendTensor.t.split(all_results_concat, M_sizes) original_n_fields = len(eval_inputs) scalar_fields = all_results_split[:original_n_fields] diff --git a/gempy_engine/modules/kernel_constructor/_covariance_assembler.py b/gempy_engine/modules/kernel_constructor/_covariance_assembler.py index ddae6a81..f928f671 100644 --- a/gempy_engine/modules/kernel_constructor/_covariance_assembler.py +++ b/gempy_engine/modules/kernel_constructor/_covariance_assembler.py @@ -98,10 +98,7 @@ def _get_cov_grad_legacy(cov_grad, dm, nugget, execution_mode: KernelExecutionMo matrix_shape = dm.hu.shape[0] LazyTensor = _lazy_tensor_class() - if BackendTensor.engine_backend == AvailableBackends.PYTORCH: - diag_ = BackendTensor.t.arange(matrix_shape).reshape(-1, 1).type(BackendTensor.dtype_obj) - else: - diag_ = np.arange(matrix_shape).reshape(-1, 1).astype(BackendTensor.dtype) + diag_ = BackendTensor.arange(matrix_shape, dtype=BackendTensor.dtype_obj).reshape(-1, 1) diag_i = LazyTensor(diag_[:, None]) diag_j = LazyTensor(diag_[None, :]) @@ -166,10 +163,7 @@ def _get_cov_surface_points_legacy( matrix_shape = k_rest_ref.shape[0] LazyTensor = _lazy_tensor_class() - if BackendTensor.engine_backend == AvailableBackends.PYTORCH: - diag_ = BackendTensor.t.arange(matrix_shape).reshape(-1, 1).type(BackendTensor.dtype_obj) - else: - diag_ = np.arange(matrix_shape).reshape(-1, 1).astype(BackendTensor.dtype) + diag_ = BackendTensor.arange(matrix_shape, dtype=BackendTensor.dtype_obj).reshape(-1, 1) nuggets = BackendTensor.t.zeros(matrix_shape, dtype=BackendTensor.dtype_obj) nuggets[grad_matrix_size:grad_matrix_size + nugget.shape[0]] += nugget @@ -268,7 +262,7 @@ def _nugget_diagonal( LazyTensor = _lazy_tensor_class() - diag_ = BackendTensor.t.arange(matrix_size, dtype=nugget.dtype).reshape(-1, 1) + diag_ = BackendTensor.arange(matrix_size, dtype=nugget.dtype).reshape(-1, 1) diag_i = LazyTensor(diag_[:, None]) diag_j = LazyTensor(diag_[None, :]) values_j = LazyTensor(values[None, :, None]) diff --git a/tests/test_common/test_api/test_concurrent_options.py b/tests/test_common/test_api/test_concurrent_options.py new file mode 100644 index 00000000..61a31eab --- /dev/null +++ b/tests/test_common/test_api/test_concurrent_options.py @@ -0,0 +1,60 @@ +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier + +import pytest + +from gempy_engine.API.model import model_api +from gempy_engine.core.data import InterpolationOptions +from gempy_engine.core.data.options.temp_interpolation_values import TempInterpolationValues +from gempy_engine.modules.evaluator.symbolic_evaluator import _validate_stacked_dimensions + + +def test_computations_do_not_share_volatile_option_state(monkeypatch): + options = InterpolationOptions.from_args( + range=1.0, + c_o=1.0, + number_octree_levels=2, + mesh_extraction=False, + ) + barrier = Barrier(2) + + monkeypatch.setattr(model_api.WeightCache, "initialize_cache_dir", lambda: None) + monkeypatch.setattr(model_api.WeightCache, "clear_cache", lambda: None) + monkeypatch.setattr(model_api.BackendTensor, "clear_gpu_memory", lambda: None) + monkeypatch.setattr(model_api, "_check_input_validity", lambda *_: None) + + def observe_options(interpolation_input, options, data_descriptor): + level = interpolation_input + options.temp_interpolation_values.current_octree_level = level + barrier.wait() + return options.temp_interpolation_values.current_octree_level + + monkeypatch.setattr(model_api, "interpolate_n_octree_levels", observe_options) + + class FakeSolutions: + def __init__(self, octrees_output, **_): + self.observed_level = octrees_output + + monkeypatch.setattr(model_api, "Solutions", FakeSolutions) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(model_api.compute_model, level, options, None) + for level in (1, 3) + ] + + assert [future.result().observed_level for future in futures] == [1, 3] + assert options.temp_interpolation_values == TempInterpolationValues() + + +def test_stacked_dimension_validation_reports_range_mismatch(): + class FakeKernel: + shape = (4, 8, 1) + + class FakeWeights: + shape = (4,) + + with pytest.raises(ValueError, match="range dimensions=\\(4, 16\\)") as error: + _validate_stacked_dimensions(FakeKernel(), FakeWeights(), M_sizes=[8, 8], N_sizes=[2, 2]) + + assert "kernel=(4, 8, 1)" in str(error.value) diff --git a/tests/test_common/test_api/test_stack_options_override.py b/tests/test_common/test_api/test_stack_options_override.py index fa2a562d..56057c08 100644 --- a/tests/test_common/test_api/test_stack_options_override.py +++ b/tests/test_common/test_api/test_stack_options_override.py @@ -86,8 +86,10 @@ def test_stack_options_override_serial(override_setup): @pytest.mark.skipif(not PYKEOPS_AVAILABLE, reason="pykeops not installed") -def test_stack_options_override_flat(override_setup): +@pytest.mark.parametrize("number_octree_levels", [1, 2]) +def test_stack_options_override_flat(override_setup, number_octree_levels): ii, global_options, ts = override_setup + global_options.evaluation_options.number_octree_levels = number_octree_levels custom_options = InterpolationOptions.from_args( range=5.0, diff --git a/tests/test_common/test_modules/test_kernel_constructor/test_nugget_propagation.py b/tests/test_common/test_modules/test_kernel_constructor/test_nugget_propagation.py index d3ac5b00..173e0b98 100644 --- a/tests/test_common/test_modules/test_kernel_constructor/test_nugget_propagation.py +++ b/tests/test_common/test_modules/test_kernel_constructor/test_nugget_propagation.py @@ -125,6 +125,36 @@ def test_surface_preprocessing_preserves_nugget_components_and_surface_ids(simpl ) +def test_surface_preprocessing_uses_backend_device_when_torch_default_differs(simple_model_2): + if BackendTensor.engine_backend is not AvailableBackends.PYTORCH or not BackendTensor.use_gpu: + pytest.skip("PyTorch GPU-only device regression test") + import torch + + previous_default_device = torch.get_default_device() + try: + torch.set_default_device("cpu") + surface_points, _, _, descriptor = deepcopy(simple_model_2) + + internal = surface_points_preprocess(surface_points, descriptor.tensors_structure) + + assert internal.surface_ids.device.type == BackendTensor.device.type + finally: + torch.set_default_device(previous_default_device) + + +def test_pytorch_repeat_moves_existing_tensors_to_backend_device(): + if BackendTensor.engine_backend is not AvailableBackends.PYTORCH or not BackendTensor.use_gpu: + pytest.skip("PyTorch GPU-only device regression test") + import torch + + values = torch.tensor([0.0], device="cpu") + repeats = torch.tensor([1], device=BackendTensor.device) + + result = BackendTensor.t.repeat(values, repeats, 0) + + assert result.device.type == BackendTensor.device.type + + def test_all_nuggets_receive_pytorch_gradients(simple_model_2): if BackendTensor.engine_backend is not AvailableBackends.PYTORCH: pytest.skip("PyTorch-only autograd test")