diff --git a/MICRO_ANISOTROPIC_FIELD_DEFORMATION.md b/MICRO_ANISOTROPIC_FIELD_DEFORMATION.md index 539c97c1..36c7b98d 100644 --- a/MICRO_ANISOTROPIC_FIELD_DEFORMATION.md +++ b/MICRO_ANISOTROPIC_FIELD_DEFORMATION.md @@ -1,6 +1,9 @@ # Micro Anisotropic Field Deformation Notes -This document summarizes the current prototype and the recommended next steps for moving the micro-correction idea into GemPy Engine with minimal disruption to the existing macro interpolation path. +This document records the historical prototype and design exploration. For the +current authored-contact API and limitations, see [docs/micro_points.md](docs/micro_points.md). +The historical covariance solve and gradient-derived metrics below now live in +`tests/test_common/test_modules/test_evaluator/micro_reference.py`, not the production API. ## Goal @@ -333,7 +336,9 @@ micro_weights = solve_micro_weights( ) ``` -10. Store on `options.evaluation_options.micro_anisotropic` and evaluate final field. +10. Prototype-only: apply fitted weights explicitly; the current authored API + uses `options.micro_options` for settings and a stack-local `MicroCorrection` + for fitted state (see [docs/micro_points.md](docs/micro_points.md)). ### 3D Frame Construction diff --git a/docs/micro_points.md b/docs/micro_points.md new file mode 100644 index 00000000..188bd819 --- /dev/null +++ b/docs/micro_points.md @@ -0,0 +1,66 @@ +# Authored micro points + +Micro contacts add a local anisotropic scalar-field correction after the ordinary +macro cokriging solve. They do not enter the macro covariance system. The feature +is opt-in: authored points are ignored with a warning while +`options.micro_options.enabled` is `False`. + +For a GemPy `GeoModel`, assign contacts to the owning structural element. Each +4-by-4 support transform has its contact position in the last column and its +three support axes in the upper-left 3-by-3 block: + +```python +import numpy as np +import gempy as gp + +support = np.eye(4)[None].copy() +support[0, :3, 3] = [0.5, 0.5, 0.52] +support[0, :3, :3] = np.diag([0.5, 0.5, 0.5]) +element = model.structural_frame.structural_groups[0].elements[0] +element.micro_points = gp.data.MicroPointsTable.from_transforms( + support, names=element.name +) +model.interpolation_options.micro_options.enabled = True +gp.compute_model(model) +``` + +Alternatively, engine callers pass `MicroPoints(points, anisotropy_matrices, +nuggets, surface_indices)` to `InterpolationInput(..., micro_points=...)`. +`surface_indices` are zero-based **global** surface indices in structural-frame +order. Inputs have shapes `(N, 3)`, `(N, 3, 3)`, `(N,)`, and `(N,)` respectively. +Internal stack subsets rebase surface indices to the owning stack. Model preparation +validates ownership and supported stack types once, then shares contact queries +across stacks; the authored arrays are not converted or mutated in place. +The matrices transform coordinate differences into each contact's local metric; +GemPy builds them from the inverse support axes after coordinate transforms. + +Configure `options.micro_options.kernel_type` (`"exponential"`, `"matern_3_2"`, +or `"matern_5_2"`), `kernel_range`, `nugget`, `strength`, and +`preserve_macro_points` as needed. With macro preservation enabled (the default), +surface points become additional zero-correction constraints. Micro weights are +fitted separately per stack at each interpolation level; neither inputs nor +options store fitted state. Set `enabled = False` to deactivate without removing +authored contacts. PyTorch fits preserve gradients through coordinates, matrices, +and contact nuggets; nearest metric selection for macro constraints is discrete. + +The final scalar field is the faulted macro field plus the stack-local micro +deformation. Faults do not mask, displace, or split the micro kernel. Contacts +on fault-affected stratigraphic stacks fit residuals against the macro field +including upstream ordinary or finite faults. When any contact's stack is enabled, +every stack evaluates the same coordinate layout: grid, all macro surface points, +then all authored micro contacts (including contacts on disabled stacks). These +last coordinates are evaluation queries only, not macro interpolation constraints +or public grid cells. The fit reads macro values directly from that evaluation; +there is no separate macro kernel call for the contacts. Upstream fault rows use +the same suffix, including activation and finite-fault projection. Fault-row +minimum and finite-fault normal selection use only the grid and macro surface +point prefix, so appended contacts cannot change the reference frame. Enabled +authored micro points on fault stacks are rejected, whether or not the fault has +upstream dependencies. Disabled authored points on fault stacks are ignored. +The micro fit uses a dense NumPy or PyTorch solve (not PyKeOps); macro PyKeOps +acceleration remains available and micro evaluation is tensor-native (the dense +fit and correction are not PyKeOps-accelerated). Flat stacks +retain fused macro scalar evaluation when possible; gradients use the non-fused +macro evaluator. Nearest-center metric assignment for preserved macro points is +discrete. Finite-fault projection uses NumPy and is not end-to-end differentiable. +Enabled authored micro contacts on external-function stacks remain unsupported. diff --git a/gempy_engine/API/interp_single/_aux_faults_ops.py b/gempy_engine/API/interp_single/_aux_faults_ops.py index e33c8300..77363326 100644 --- a/gempy_engine/API/interp_single/_aux_faults_ops.py +++ b/gempy_engine/API/interp_single/_aux_faults_ops.py @@ -11,12 +11,10 @@ def _grab_stack_fault_data(_all_stack_values_block, _interpolation_input_i, _stack_structure, grid_size:int) -> FaultsData: fault_data = _interpolation_input_i.fault_values or FaultsData() fault_data.fault_values_everywhere = _all_stack_values_block[_stack_structure.active_faults_relations] - fv_on_all_sp = fault_data.fault_values_everywhere[:, grid_size:] + fv_on_all_sp = fault_data.fault_values_everywhere[:, grid_size:_interpolation_input_i.macro_reference_size] fault_data.fault_values_on_sp = fv_on_all_sp[:, _interpolation_input_i.slice_feature] return fault_data - - def _options_with_finite_fault_gradients( options: InterpolationOptions, fault_input: FaultsData | None, @@ -35,7 +33,8 @@ def _modify_faults_values_output( xyz_to_interpolate: np.ndarray, ) -> np.ndarray: values_on_all_xyz = output.values_on_all_xyz - val_min = BackendTensor.t.min(values_on_all_xyz, axis=1).reshape(-1, 1) # ? Is this as good as it gets? + reference_size = output.exported_fields._macro_reference_size or len(xyz_to_interpolate) + val_min = BackendTensor.t.min(values_on_all_xyz[:, :reference_size], axis=1).reshape(-1, 1) shifted_vals = (values_on_all_xyz - val_min) # * Shift values between 0 and 1... hopefully if not fault_input.finite_fault_defined: return shifted_vals @@ -65,12 +64,12 @@ def _modify_faults_values_output( ) gradient_matrix = np.stack(gradients_np, axis=-1) - valid_gradient = np.linalg.norm(gradient_matrix, axis=1) > 1e-12 + valid_gradient = np.linalg.norm(gradient_matrix[:reference_size], axis=1) > 1e-12 if not np.any(valid_gradient): raise ValueError("Cannot determine finite-fault frame from near-zero gradients") center = np.asarray(fault_input.finite_fault.center) - distances_to_center = np.linalg.norm(projected_points - center, axis=1) + distances_to_center = np.linalg.norm(projected_points[:reference_size] - center, axis=1) center_index = np.argmin(np.where(valid_gradient, distances_to_center, np.inf)) finite_fault_scalar_np = fault_input.finite_fault.calculate_slip( points=projected_points, diff --git a/gempy_engine/API/interp_single/_interp_scalar_field.py b/gempy_engine/API/interp_single/_interp_scalar_field.py index 012cc04c..6ddc48bb 100644 --- a/gempy_engine/API/interp_single/_interp_scalar_field.py +++ b/gempy_engine/API/interp_single/_interp_scalar_field.py @@ -102,7 +102,7 @@ def _solve_interpolation_result( def _evaluate_sys_eq(eval_input: Union[SolverInput, EvaluatorInput], weights: np.ndarray, options: InterpolationOptions, - grid: EngineGrid | None = None) -> ExportedFields: + grid: EngineGrid | None = None) -> ExportedFields: inverse = None if options.evaluation_options.deduplicate_octree_corners: eval_input, inverse = _deduplicate_corners(eval_input, grid) diff --git a/gempy_engine/API/interp_single/_interp_single_feature.py b/gempy_engine/API/interp_single/_interp_single_feature.py index ce7d32c2..554f275c 100644 --- a/gempy_engine/API/interp_single/_interp_single_feature.py +++ b/gempy_engine/API/interp_single/_interp_single_feature.py @@ -16,8 +16,10 @@ from ...core.data.kernel_classes.faults import FaultsData from ...core.data.options import InterpolationOptions from ...core.data.scalar_field_output import ScalarFieldOutput +from ...core.data.stack_relation_type import StackRelationType from ...modules.activator import activator_interface from ...modules.data_preprocess import data_preprocess_interface +from ...modules.evaluator.micro_correction import apply_micro_correction, fit_micro_correction def interpolate_feature_with_cokrig(interpolation_input: InterpolationInput, @@ -36,11 +38,15 @@ def interpolate_feature_with_cokrig(interpolation_input: InterpolationInput, weights = compute_weights(solver_input, stack_number, options) exported_fields: ExportedFields = _evaluate_sys_eq(solver_input, weights, options, grid=grid) + correction = fit_micro_correction(interpolation_input, exported_fields.scalar_field_everywhere, options, + data_shape.number_of_points_per_surface) + apply_micro_correction(exported_fields, xyz, correction) exported_fields.set_structure_values( reference_sp_position=data_shape.reference_sp_position, slice_feature=interpolation_input.slice_feature, - grid_size=interpolation_input.grid.len_all_grids + grid_size=interpolation_input.grid.len_all_grids, + macro_reference_size=interpolation_input.macro_reference_size ) exported_fields.debug = solver_input.debug @@ -64,7 +70,11 @@ def interpolate_feature_with_external_function(interpolation_input: Interpolatio grid = interpolation_input.grid # region Interpolate scalar field - xyz = grid.values + shared_queries = (interpolation_input.evaluation_micro_points is not None or + interpolation_input.stack_relation is StackRelationType.FAULT) + xyz = (data_preprocess_interface.prepare_grid(grid.values, interpolation_input.all_surface_points, + interpolation_input.evaluation_micro_points) + if shared_queries else grid.values) exported_fields: ExportedFields = _interpolate_external_function( interp_funct=external_interp_funct, @@ -74,7 +84,8 @@ def interpolate_feature_with_external_function(interpolation_input: Interpolatio exported_fields.set_structure_values( reference_sp_position=None, slice_feature=None, - grid_size=xyz.shape[0] + grid_size=grid.len_all_grids, + macro_reference_size=interpolation_input.macro_reference_size if shared_queries else None ) output = _segment(exported_fields, external_segment_funct, grid, interpolation_input, options, xyz) @@ -92,7 +103,8 @@ def input_preprocess(data_shape: TensorsStructure, interpolation_input: Interpol # * We need to interpolate in ALL the surface points not only the surface points of the stack grid_internal: np.ndarray = data_preprocess_interface.prepare_grid( grid=grid.values, - surface_points=interpolation_input.all_surface_points + surface_points=interpolation_input.all_surface_points, + micro_points=interpolation_input.evaluation_micro_points ) fault_values: FaultsData = interpolation_input.fault_values @@ -149,11 +161,14 @@ def _scalar_field_segmentation(exported_fields: ExportedFields, external_segment def _interpolate_external_function(interp_funct, xyz): + def field(func): + return BackendTensor.t.array(func(xyz), dtype=BackendTensor.dtype) if func is not None else None + exported_fields = ExportedFields( - _scalar_field=interp_funct.implicit_function(xyz), - _gx_field=interp_funct.gx_function(xyz) if interp_funct.gx_function is not None else None, - _gy_field=interp_funct.gy_function(xyz) if interp_funct.gy_function is not None else None, - _gz_field=interp_funct.gz_function(xyz) if interp_funct.gz_function is not None else None, + _scalar_field=field(interp_funct.implicit_function), + _gx_field=field(interp_funct.gx_function), + _gy_field=field(interp_funct.gy_function), + _gz_field=field(interp_funct.gz_function), _scalar_field_at_surface_points=interp_funct.scalar_field_at_surface_points ) return exported_fields diff --git a/gempy_engine/API/interp_single/_multi_scalar_field_manager.py b/gempy_engine/API/interp_single/_multi_scalar_field_manager.py index 0dfd820a..3880099e 100644 --- a/gempy_engine/API/interp_single/_multi_scalar_field_manager.py +++ b/gempy_engine/API/interp_single/_multi_scalar_field_manager.py @@ -19,6 +19,7 @@ from ...core.data.scalar_field_output import ScalarFieldOutput from ...core.data.stack_relation_type import StackRelationType from ...core.data.stacks_structure import StacksStructure +from ...modules.data_preprocess.data_preprocess_interface import prepare_grid # @off @@ -33,7 +34,7 @@ def interpolate_all_fields(interpolation_input: InterpolationInput, options: Int ) if (os.getenv("GEMPY_FLAT_STACKS", "False").lower() in ("true", "1", "t", "y", "yes") and - BackendTensor.use_pykeops + BackendTensor.use_pykeops # and not has_external_functions ): all_scalar_fields_outputs: List[ScalarFieldOutput] = _interpolate_stack_flat(data_descriptor, interpolation_input, options) @@ -57,7 +58,7 @@ def _interpolate_stack_flat(root_data_descriptor: InputDataDescriptor, root_inte options: InterpolationOptions) -> ScalarFieldOutput | List[ScalarFieldOutput]: stack_structure = root_data_descriptor.stack_structure - xyz_to_interpolate_size: int = root_interpolation_input.grid.len_all_grids + root_interpolation_input.surface_points.n_points + xyz_to_interpolate_size: int = root_interpolation_input.evaluation_size all_stack_values_block: np.ndarray = BackendTensor.t.zeros( (stack_structure.n_stacks, xyz_to_interpolate_size), dtype=BackendTensor.dtype_obj @@ -95,7 +96,7 @@ def _interpolate_stack(root_data_descriptor: InputDataDescriptor, root_interpola all_scalar_fields_outputs: List[ScalarFieldOutput | None] = [None] * stack_structure.n_stacks - xyz_to_interpolate_size: int = root_interpolation_input.grid.len_all_grids + root_interpolation_input.surface_points.n_points + xyz_to_interpolate_size: int = root_interpolation_input.evaluation_size all_stack_values_block: np.ndarray = BackendTensor.t.zeros( (stack_structure.n_stacks, xyz_to_interpolate_size), dtype=BackendTensor.dtype_obj) # * Used for faults @@ -113,6 +114,7 @@ def _interpolate_stack(root_data_descriptor: InputDataDescriptor, root_interpola all_interpolation_input=root_interpolation_input, stack_structure=stack_structure ) + fault_input = interpolation_input_i.fault_values output: ScalarFieldOutput = interpolate_feature_with_external_function( interpolation_input=interpolation_input_i, @@ -154,7 +156,10 @@ def _interpolate_stack(root_data_descriptor: InputDataDescriptor, root_interpola values_output = _modify_faults_values_output( # ! This is all_STACK_values_block (not all_scalar_fields_outputs) fault_input=fault_input, output=output, - xyz_to_interpolate=solver_input.xyz_to_interpolate + xyz_to_interpolate=(solver_input.xyz_to_interpolate if stack_structure.interp_function is None else + prepare_grid(interpolation_input_i.grid.values, + interpolation_input_i.all_surface_points, + interpolation_input_i.evaluation_micro_points)) ) all_stack_values_block[i, :] = values_output diff --git a/gempy_engine/API/interp_single/_stack_ops.py b/gempy_engine/API/interp_single/_stack_ops.py index 4a5c5b01..fc504db1 100644 --- a/gempy_engine/API/interp_single/_stack_ops.py +++ b/gempy_engine/API/interp_single/_stack_ops.py @@ -22,6 +22,7 @@ from ...modules.activator import activator_interface from ...modules.data_preprocess import data_preprocess_interface +from ...modules.evaluator.micro_correction import apply_micro_correction, fit_micro_correction @dataclass @@ -91,7 +92,9 @@ def process_chunk(state: InterpolationState, chunk: list[int]): state.solver_inputs[i] = chunk_solver_inputs[idx] # Evaluate this chunk - if any(fault_input.finite_fault_defined for fault_input in (item.fault_values for item in chunk_interpolation_inputs)): + # The fused PyKeOps evaluator cannot concatenate scalar and gradient lazy kernels. + if (any(item.fault_values.finite_fault_defined for item in chunk_interpolation_inputs) + or any(opt.compute_scalar_gradient for opt in chunk_options)): chunk_eval_inputs, chunk_exported_fields = _evaluate( interpolation_inputs=chunk_interpolation_inputs, options=state.options, @@ -114,6 +117,10 @@ def process_chunk(state: InterpolationState, chunk: list[int]): for idx, i in enumerate(chunk): state.eval_inputs[i] = chunk_eval_inputs[idx] + correction = fit_micro_correction(chunk_interpolation_inputs[idx], + chunk_exported_fields[idx].scalar_field_everywhere, + chunk_options[idx], chunk_tensor_structs[idx].number_of_points_per_surface) + apply_micro_correction(chunk_exported_fields[idx], chunk_eval_inputs[idx].xyz_to_interpolate, correction) # Segment this chunk chunk_outputs = _segment( @@ -238,8 +245,8 @@ def _evaluate(interpolation_inputs: list[InterpolationInput], options: Interpola def _evaluate_optimized(interpolation_inputs: list[InterpolationInput], options: InterpolationOptions, solver_inputs, stack_structure: StacksStructure, - tensor_structs: list[TensorsStructure], stack_indices: list[int] | None = None, - options_per_stack: list[InterpolationOptions] | None = None) -> tuple[list[EvaluatorInput], list[ExportedFields]]: + tensor_structs: list[TensorsStructure], stack_indices: list[int] | None = None, + options_per_stack: list[InterpolationOptions] | None = None) -> tuple[list[EvaluatorInput], list[ExportedFields]]: from gempy_engine.modules.evaluator.symbolic_evaluator import symbolic_evaluator_optimized_stacked eval_inputs: list[EvaluatorInput] = [] @@ -270,7 +277,7 @@ def _evaluate_optimized(interpolation_inputs: list[InterpolationInput], options: exported_fields_list: list[ExportedFields] = symbolic_evaluator_optimized_stacked( eval_inputs=reduced_inputs, weights_list=weights_list, - options_list=options_list + options_list=options_list, ) for idx, exported_fields in enumerate(exported_fields_list): @@ -379,7 +386,6 @@ def _process_external_chunk(state: InterpolationState, chunk: list[int]): all_interpolation_input=state.root_interpolation_input, stack_structure=state.stack_structure ) - output: ScalarFieldOutput = interpolate_feature_with_external_function( interpolation_input=interpolation_input_i, options=options_i, @@ -388,3 +394,9 @@ def _process_external_chunk(state: InterpolationState, chunk: list[int]): ) state.all_scalar_fields_outputs[i] = output + if interpolation_input_i.stack_relation is StackRelationType.FAULT: + xyz = data_preprocess_interface.prepare_grid( + interpolation_input_i.grid.values, interpolation_input_i.all_surface_points, + interpolation_input_i.evaluation_micro_points) + state.all_stack_values_block[i, :] = _modify_faults_values_output( + interpolation_input_i.fault_values, output, xyz) diff --git a/gempy_engine/API/model/model_api.py b/gempy_engine/API/model/model_api.py index 89cb670a..c6763e46 100644 --- a/gempy_engine/API/model/model_api.py +++ b/gempy_engine/API/model/model_api.py @@ -20,6 +20,7 @@ from ...core.utils import gempy_profiler_decorator from ...core.exceptions import GemPyEngineInputError from ...core.data.options.temp_interpolation_values import TempInterpolationValues +from ...modules.data_preprocess.micro_points import prepare_micro_points 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 @@ -43,6 +44,8 @@ def compute_model(interpolation_input: InterpolationInput, options: Interpolatio # Check input is valid _check_input_validity(interpolation_input, options, data_descriptor) # TODO + prepare_micro_points(interpolation_input, options, data_descriptor.stack_structure) + output: list[OctreeLevel] = interpolate_n_octree_levels( interpolation_input=interpolation_input, options=options, diff --git a/gempy_engine/core/data/exported_fields.py b/gempy_engine/core/data/exported_fields.py index 5a653d5e..195ad915 100644 --- a/gempy_engine/core/data/exported_fields.py +++ b/gempy_engine/core/data/exported_fields.py @@ -17,6 +17,7 @@ class ExportedFields: _n_points_per_surface: Optional[np.ndarray] = None _slice_feature: Optional[slice] = field(default_factory=lambda: slice(None, None)) # Slice all the surface points _grid_size: Optional[int] = None + _macro_reference_size: Optional[int] = None _scalar_field_at_surface_points: Optional[np.ndarray] = None scalar_field_at_fault_shell: Optional[np.ndarray] = None @@ -24,15 +25,17 @@ class ExportedFields: debug = None def set_structure_values(self, reference_sp_position: Optional[np.ndarray], slice_feature: Optional[slice], - grid_size: int): + grid_size: int, macro_reference_size: int | None = None): self._n_points_per_surface = reference_sp_position self._grid_size = grid_size self._slice_feature = slice_feature + self._macro_reference_size = macro_reference_size def set_structure_values_from_eval_input(self, evaluator_input: EvaluatorInput): self._n_points_per_surface = evaluator_input._n_points_per_surface self._grid_size = evaluator_input._grid_size self._slice_feature = evaluator_input._slice_feature + self._macro_reference_size = evaluator_input._macro_reference_size @property def n_points_per_surface(self) -> np.ndarray: @@ -53,7 +56,7 @@ def scalar_field_at_surface_points(self) -> Optional[np.ndarray]: if self.scalar_field_at_fault_shell is not None: # * For now this has priority over everything else return self.scalar_field_at_fault_shell elif self._scalar_field_at_surface_points is None: - scalar_field_at_all_sp = self._scalar_field[self.grid_size:] + scalar_field_at_all_sp = self._scalar_field[self.grid_size:self._macro_reference_size] scalar_field_at_feature_sp = scalar_field_at_all_sp[self._slice_feature] scalar_field_at_one_point_per_surface = scalar_field_at_feature_sp[self.npf] return scalar_field_at_one_point_per_surface @@ -66,7 +69,7 @@ def scalar_field_at_surface_points(self, value): @property def scalar_field(self): - if self._slice_feature is None: + if self._grid_size is None: return self._scalar_field return self._scalar_field[:self._grid_size] diff --git a/gempy_engine/core/data/internal_structs.py b/gempy_engine/core/data/internal_structs.py index 6d1149a8..59007535 100644 --- a/gempy_engine/core/data/internal_structs.py +++ b/gempy_engine/core/data/internal_structs.py @@ -99,6 +99,7 @@ class EvaluatorInput: _n_points_per_surface: Optional[np.ndarray] = None _slice_feature: Optional[slice] = field(default_factory=lambda: slice(None, None)) # Slice all the surface points _grid_size: Optional[int] = None + _macro_reference_size: Optional[int] = None def __init__(self, solver_input: SolverInput_v2, @@ -109,6 +110,7 @@ def __init__(self, self.solver_input = solver_input self._n_points_per_surface = tensor_struct.reference_sp_position self._slice_feature = interpolation_input.slice_feature + self._macro_reference_size = interpolation_input.macro_reference_size if only_surface_points: xyz_to_interpolate = interpolation_input.all_surface_points.sp_coords @@ -116,7 +118,8 @@ def __init__(self, else: xyz_to_interpolate: np.ndarray = data_preprocess_interface.prepare_grid( grid=interpolation_input.grid.values, - surface_points=interpolation_input.all_surface_points + surface_points=interpolation_input.all_surface_points, + micro_points=interpolation_input.evaluation_micro_points ) self._grid_size = interpolation_input.grid.len_all_grids diff --git a/gempy_engine/core/data/interpolation_input.py b/gempy_engine/core/data/interpolation_input.py index f9d9c1be..883715ad 100644 --- a/gempy_engine/core/data/interpolation_input.py +++ b/gempy_engine/core/data/interpolation_input.py @@ -13,6 +13,7 @@ from .stacks_structure import StacksStructure from .kernel_classes.faults import FaultsData from .kernel_classes.server.input_parser import InterpolationInputSchema +from .micro_points import MicroPoints @dataclass @@ -28,6 +29,10 @@ class InterpolationInput: segmentation_function: Optional[callable] = None # * From scalar field to values _all_surface_points: SurfacePoints = None + micro_points: Optional[MicroPoints] = None + # Prepared by prepare_micro_points: shared queries and this stack's rows in them. + _all_micro_points: Optional[MicroPoints] = None + _micro_indices: Optional[np.ndarray] = None # region per model ? Not sure what I mean here @@ -38,7 +43,8 @@ class InterpolationInput: def __init__(self, surface_points: SurfacePoints, orientations: Orientations, grid: EngineGrid, unit_values: Optional[np.ndarray] = None, segmentation_function: Optional[callable] = None, - stack_relation: StackRelationType = StackRelationType.ERODE, weights: list[np.ndarray] = None): + stack_relation: StackRelationType = StackRelationType.ERODE, weights: list[np.ndarray] = None, + micro_points: Optional[MicroPoints] = None): if weights is None: weights = [] @@ -50,6 +56,7 @@ def __init__(self, surface_points: SurfacePoints, orientations: Orientations, gr self.segmentation_function = segmentation_function self.stack_relation = stack_relation self.weights = weights + self.micro_points = micro_points # @ on @@ -76,6 +83,13 @@ def from_interpolation_input_subset(cls, all_interpolation_input: "Interpolation unit_values = all_interpolation_input.unit_values[cum_number_surfaces_l0:cum_number_surfaces_l1] grid = all_interpolation_input.grid + micro = all_interpolation_input.micro_points + if micro is not None: + start = int(stack_structure.number_of_surfaces_per_stack[:stack_number].sum()) + stop = start + int(stack_structure.number_of_surfaces_per_stack[stack_number]) + mask = (micro.surface_indices >= start) & (micro.surface_indices < stop) + micro = MicroPoints(micro.points[mask], micro.anisotropy_matrices[mask], + micro.nuggets[mask], micro.surface_indices[mask] - start) # * (miguel 24) This interpolation input goes on the InterpOutput so we are not computing its gradients ii_subset: InterpolationInput = cls( @@ -84,12 +98,15 @@ def from_interpolation_input_subset(cls, all_interpolation_input: "Interpolation grid=grid, unit_values=unit_values, stack_relation=stack_structure.active_masking_descriptor, - weights=(all_interpolation_input.weights[stack_number] if stack_number < len(all_interpolation_input.weights) else None) + weights=(all_interpolation_input.weights[stack_number] if stack_number < len(all_interpolation_input.weights) else None), + micro_points=micro, ) # ! Setting this on the constructor does not work with data classes. ii_subset.fault_values = stack_structure.active_faults_input_data ii_subset.all_surface_points = all_interpolation_input.surface_points + ii_subset._all_micro_points = all_interpolation_input._all_micro_points + ii_subset._micro_indices = np.flatnonzero(mask) if all_interpolation_input._all_micro_points is not None else None return ii_subset @@ -127,6 +144,27 @@ def set_temp_grid(self, value): def slice_feature(self): return self.surface_points.slice_feature + @property + def evaluation_micro_points(self): + return self._all_micro_points + + @property + def micro_indices(self): + return self._micro_indices + + @property + def macro_reference_size(self): + return self.grid.len_all_grids + self.all_surface_points.n_points + + @property + def micro_slice(self): + start = self.macro_reference_size + return slice(start, start + (len(self._all_micro_points.points) if self._all_micro_points is not None else 0)) + + @property + def evaluation_size(self): + return self.micro_slice.stop + @property def fault_values(self): if self._fault_values is None: diff --git a/gempy_engine/core/data/micro_points.py b/gempy_engine/core/data/micro_points.py new file mode 100644 index 00000000..bd79b4fa --- /dev/null +++ b/gempy_engine/core/data/micro_points.py @@ -0,0 +1,39 @@ +from dataclasses import dataclass + +import numpy as np +from gempy_engine.config import is_pytorch_installed + +if is_pytorch_installed: + import torch + MicroArray = np.ndarray | torch.Tensor +else: + MicroArray = np.ndarray + + +@dataclass +class MicroPoints: + """Authored contacts indexed by surface in the owning InterpolationInput. + + Root inputs use global structural-frame indices; stack subsets use local indices. + """ + points: MicroArray + anisotropy_matrices: MicroArray + nuggets: MicroArray + surface_indices: np.ndarray + + def __post_init__(self): + indices = np.asarray(self.surface_indices) + n = len(self.points) + if (self.points.shape != (n, 3) or self.anisotropy_matrices.shape != (n, 3, 3) + or self.nuggets.shape != (n,) or indices.shape != (n,) + or not np.issubdtype(indices.dtype, np.integer)): + raise ValueError("Invalid micro_points shapes or surface_indices dtype") + for name in ("points", "anisotropy_matrices", "nuggets"): + values = getattr(self, name) + finite = (torch.isfinite(values).all() if is_pytorch_installed and isinstance(values, torch.Tensor) + else np.isfinite(values).all()) + if not finite: + raise ValueError("micro_points must be finite") + if (self.nuggets < 0).any() or (indices < 0).any(): + raise ValueError("micro_points must have nonnegative nuggets and indices") + self.surface_indices = indices.astype(np.int64) diff --git a/gempy_engine/core/data/options/evaluation_options.py b/gempy_engine/core/data/options/evaluation_options.py index 6ed35f60..5da95ca4 100644 --- a/gempy_engine/core/data/options/evaluation_options.py +++ b/gempy_engine/core/data/options/evaluation_options.py @@ -5,7 +5,6 @@ from typing_extensions import deprecated -from .micro_anisotropic_options import MicroAnisotropicOptions class OctreeRefinementMode(str, enum.Enum): @@ -50,8 +49,6 @@ class EvaluationOptions: evaluation_chunk_size: int = 500_000 - micro_anisotropic: MicroAnisotropicOptions = field(default_factory=MicroAnisotropicOptions) - compute_scalar: bool = True compute_scalar_gradient: bool = False diff --git a/gempy_engine/core/data/options/interpolation_options.py b/gempy_engine/core/data/options/interpolation_options.py index df933e30..4b32a7e7 100644 --- a/gempy_engine/core/data/options/interpolation_options.py +++ b/gempy_engine/core/data/options/interpolation_options.py @@ -5,6 +5,7 @@ import gempy_engine.config from .evaluation_options import MeshExtractionMaskingOptions, EvaluationOptions +from .micro_anisotropic_options import MicroAnisotropicOptions from .temp_interpolation_values import TempInterpolationValues from ..kernel_classes.kernel_functions import AvailableKernelFunctions from .kernel_options import KernelOptions, NuggetImplementation @@ -31,6 +32,7 @@ class CacheMode(enum.Enum): # @off kernel_options: KernelOptions = Field(init=True, exclude=False) # * This is the compression of the fields above and the way to go in the future evaluation_options: EvaluationOptions = Field(init=True, exclude= False) + micro_options: MicroAnisotropicOptions = Field(default_factory=MicroAnisotropicOptions) debug: bool cache_mode: CacheMode diff --git a/gempy_engine/core/data/options/micro_anisotropic_options.py b/gempy_engine/core/data/options/micro_anisotropic_options.py index 805e8bcf..2ca45fb0 100644 --- a/gempy_engine/core/data/options/micro_anisotropic_options.py +++ b/gempy_engine/core/data/options/micro_anisotropic_options.py @@ -1,19 +1,12 @@ -from typing import Literal, Optional +from typing import Literal -import numpy as np -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel MicroKernelType = Literal["exponential", "matern_3_2", "matern_5_2"] class MicroAnisotropicOptions(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - enabled: bool = False - points: Optional[np.ndarray] = None # (N, 3) micro constraint points - residuals: Optional[np.ndarray] = None # (N,) target residual values - anisotropy_matrices: Optional[np.ndarray] = None # (N, 3, 3) per-point anisotropy transforms - weights: Optional[np.ndarray] = None # (N,) solved micro weights kernel_range: float = 1.0 # range for the micro kernel kernel_type: MicroKernelType = "matern_5_2" # kernel function for micro solve + eval nugget: float = 0.0 # diagonal nugget for the micro solve diff --git a/gempy_engine/modules/data_preprocess/data_preprocess_interface.py b/gempy_engine/modules/data_preprocess/data_preprocess_interface.py index 842bbec8..7711a5e7 100644 --- a/gempy_engine/modules/data_preprocess/data_preprocess_interface.py +++ b/gempy_engine/modules/data_preprocess/data_preprocess_interface.py @@ -18,9 +18,12 @@ def prepare_orientations(orientations: Orientations) -> OrientationsInternals: return orientations_preprocess(orientations) -def prepare_grid(grid: np.ndarray, surface_points: SurfacePoints) -> np.ndarray: +def prepare_grid(grid: np.ndarray, surface_points: SurfacePoints, micro_points=None) -> np.ndarray: grid = BackendTensor.t.array(grid, dtype=BackendTensor.dtype) - concat = bt.tfnp.concatenate([grid, surface_points.sp_coords]) + arrays = [grid, surface_points.sp_coords] + if micro_points is not None: + arrays.append(BackendTensor.t.array(micro_points.points, dtype=BackendTensor.dtype)) + concat = bt.tfnp.concatenate(arrays) return concat diff --git a/gempy_engine/modules/data_preprocess/micro_points.py b/gempy_engine/modules/data_preprocess/micro_points.py new file mode 100644 index 00000000..4f5abb9d --- /dev/null +++ b/gempy_engine/modules/data_preprocess/micro_points.py @@ -0,0 +1,44 @@ +import warnings + +import numpy as np + +from ...core.data.interpolation_input import InterpolationInput +from ...core.data.options import InterpolationOptions +from ...core.data.stack_relation_type import StackRelationType +from ...core.data.stacks_structure import StacksStructure + + +def prepare_micro_points(interpolation_input: InterpolationInput, options: InterpolationOptions, + stack_structure: StacksStructure) -> None: + """Validate authored contacts and select the shared micro evaluation suffix.""" + micro = interpolation_input.micro_points + interpolation_input._all_micro_points = None + if micro is None or not len(micro.points): + return + + surface_boundaries = np.cumsum(stack_structure.number_of_surfaces_per_stack) + if (micro.surface_indices >= surface_boundaries[-1]).any(): + raise ValueError("micro_points.surface_indices contains an unknown global surface index") + enabled = False + overrides = stack_structure.interpolation_options_per_stack + for stack_index in np.unique(np.searchsorted(surface_boundaries, micro.surface_indices, side="right")): + stack_options = overrides[stack_index] if overrides is not None and overrides[stack_index] is not None else options + if stack_options.micro_options.enabled: + if stack_structure.masking_descriptor[stack_index] is StackRelationType.FAULT: + raise ValueError( + f"Stack {stack_index} is a fault stack: enabled authored micro points on fault surfaces " + "are not supported. Disable micro_options for this stack or remove its micro points." + ) + functions = stack_structure.interp_functions_per_stack + if functions is not None and functions[stack_index] is not None: + raise NotImplementedError("Authored micro points on external-function stacks are not supported") + enabled = True + continue + warnings.warn( + f"Stack {stack_index} contains micro points, but micro_options.enabled is False; " + "these points will be ignored. Set it to True to apply the correction.", + UserWarning, + stacklevel=3, + ) + if enabled: + interpolation_input._all_micro_points = micro diff --git a/gempy_engine/modules/evaluator/generic_evaluator.py b/gempy_engine/modules/evaluator/generic_evaluator.py index 13f8768f..532ab93d 100644 --- a/gempy_engine/modules/evaluator/generic_evaluator.py +++ b/gempy_engine/modules/evaluator/generic_evaluator.py @@ -67,21 +67,6 @@ def generic_evaluator( if n_chunks > 5: print(f"Chunking done: {n_chunks} chunks") - micro = options.evaluation_options.micro_anisotropic - if micro.enabled and micro.weights is not None and micro.points is not None and micro.anisotropy_matrices is not None: - from gempy_engine.modules.evaluator.micro_anisotropic_evaluator import evaluate_micro_correction - if BackendTensor.engine_backend != gempy_engine.config.AvailableBackends.numpy: - scalar_field = BackendTensor.t.to_numpy(scalar_field) - correction = evaluate_micro_correction( - xyz_to_interpolate=solver_input.xyz_to_interpolate, - micro_points=micro.points, - micro_weights=micro.weights, - anisotropy_matrices=micro.anisotropy_matrices, - kernel_range=micro.kernel_range, - kernel_type=micro.kernel_type, - ) - scalar_field = scalar_field + correction - return ExportedFields(scalar_field, gx_field, gy_field, gz_field) diff --git a/gempy_engine/modules/evaluator/micro_anisotropic_evaluator.py b/gempy_engine/modules/evaluator/micro_anisotropic_evaluator.py index 80f63265..43468078 100644 --- a/gempy_engine/modules/evaluator/micro_anisotropic_evaluator.py +++ b/gempy_engine/modules/evaluator/micro_anisotropic_evaluator.py @@ -1,7 +1,5 @@ import numpy as np -from typing import Literal, Optional - -MicroKernelType = Literal["exponential", "matern_3_2", "matern_5_2"] +from ...core.data.options.micro_anisotropic_options import MicroKernelType def _kernel_value(r: np.ndarray, kernel_type: MicroKernelType) -> np.ndarray: @@ -13,14 +11,19 @@ def _kernel_value(r: np.ndarray, kernel_type: MicroKernelType) -> np.ndarray: matern_3_2 — Matérn 3/2: K(r) = (1 + sqrt(3) r) exp(-sqrt(3) r) matern_5_2 — Matérn 5/2: K(r) = (1 + sqrt(5) r + 5r²/3) exp(-sqrt(5) r) """ + if not isinstance(r, (np.ndarray, np.generic, float, int)): + import torch + exp = torch.exp + else: + exp = np.exp if kernel_type == "exponential": - return np.exp(-r) + return exp(-r) elif kernel_type == "matern_3_2": a = np.sqrt(3.0) * r - return (1.0 + a) * np.exp(-a) + return (1.0 + a) * exp(-a) elif kernel_type == "matern_5_2": a = np.sqrt(5.0) * r - return (1.0 + a + (5.0 / 3.0) * r * r) * np.exp(-a) + return (1.0 + a + (5.0 / 3.0) * r * r) * exp(-a) else: raise ValueError(f"Unknown micro kernel type: {kernel_type}") @@ -37,153 +40,99 @@ def evaluate_micro_correction( V(x) = sum_i w_i * K(||A_i (x - p_i)|| / range) """ - M = xyz_to_interpolate.shape[0] - N = micro_points.shape[0] - correction = np.zeros(M, dtype=np.float64) - - for j in range(N): - Aj = anisotropy_matrices[j] - wj = micro_weights[j] - pj = micro_points[j] - diffs = xyz_to_interpolate - pj[np.newaxis, :] - transformed = np.einsum('ij,mj->mi', Aj, diffs) - dists = np.linalg.norm(transformed, axis=1) - r = dists / kernel_range - correction += wj * _kernel_value(r, kernel_type) - - return correction - - -def build_micro_covariance( - micro_points: np.ndarray, # (N, 3) - anisotropy_matrices: np.ndarray, # (N, 3, 3) - kernel_range: float = 1.0, - kernel_type: MicroKernelType = "exponential", - nugget: float = 0.0, -) -> np.ndarray: - """Build the symmetric NxN covariance matrix for the micro solve. - - K[i,j] = K(||A_i (p_i - p_j)|| / range) - - where K is the selected micro kernel and distances use the symmetric - metric M_ij = (A_i^T A_i + A_j^T A_j) / 2. - """ - N = micro_points.shape[0] - K = np.zeros((N, N), dtype=np.float64) - - ATA = np.einsum('nki,nkj->nij', anisotropy_matrices, anisotropy_matrices) - - for i in range(N): - for j in range(i, N): - M_ij = 0.5 * (ATA[i] + ATA[j]) - diff = micro_points[i] - micro_points[j] - dist_sq = diff @ M_ij @ diff - dist = np.sqrt(max(dist_sq, 0.0)) - r = dist / kernel_range - val = float(_kernel_value(np.array(r), kernel_type)) - K[i, j] = val - K[j, i] = val - - if nugget > 0: - np.fill_diagonal(K, K.diagonal() + nugget) - - return K - - -def solve_micro_weights( - micro_points: np.ndarray, # (N, 3) - residuals: np.ndarray, # (N,) - anisotropy_matrices: np.ndarray, # (N, 3, 3) - kernel_range: float = 1.0, - kernel_type: MicroKernelType = "exponential", - nugget: float = 0.0, -) -> np.ndarray: - """Solve K @ w = residuals for the micro correction weights. - - Returns weights array of shape (N,). - """ - K = build_micro_covariance(micro_points, anisotropy_matrices, kernel_range, kernel_type, nugget) - weights = np.linalg.solve(K, residuals) - return weights - - -def compute_macro_values_at_micro_points( - xyz_to_interpolate: np.ndarray, - weights: np.ndarray, - solver_input: 'SolverInput', - options: 'InterpolationOptions', -) -> np.ndarray: - """Extract the macro scalar field at micro point locations. - - This evaluates the macro interpolation exactly at the micro contact points - to compute residuals = target_values - macro_values. - """ - from gempy_engine.modules.evaluator.symbolic_evaluator import symbolic_evaluator - from gempy_engine.core.data.internal_structs import SolverInput - - proxy_input = SolverInput( - sp_internal=solver_input.sp_internal, - ori_internal=solver_input.ori_internal, - xyz_to_interpolate=xyz_to_interpolate, - fault_internal=solver_input._fault_internal, - ) - - exported = symbolic_evaluator(proxy_input, weights, options) - return exported.scalar_field - - -def compute_anisotropy_matrices_from_gradients( - micro_points: np.ndarray, # (N, D) - gradients: np.ndarray, # (N, D) gradient vectors - r_vertical: float = 1.0, - r_lateral: float = 10.0, -) -> np.ndarray: - """Build per-point anisotropy matrices from macro gradient directions. - - A_i = S * R_i^T - - R_i^T projects world coordinates into a local frame aligned with the gradient - (last axis = gradient direction = stratigraphic up). - S = diag(lateral scale repeated, vertical scale) - - Works for 2D and 3D. - """ - N, D = micro_points.shape - assert gradients.shape == (N, D), f"gradients shape {gradients.shape} != (N, D) {(N, D)}" - - if D == 2: - lateral_scales = np.array([1.0 / r_lateral], dtype=np.float64) - scales = np.concatenate([lateral_scales, [1.0 / r_vertical]]) - S = np.diag(scales) - else: - S = np.diag(np.array([1.0 / r_lateral, 1.0 / r_lateral, 1.0 / r_vertical])) - - matrices = np.zeros((N, D, D), dtype=np.float64) - - for i in range(N): - grad = gradients[i].astype(np.float64) - grad_norm = np.linalg.norm(grad) - if grad_norm < 1e-10: - grad = np.zeros(D, dtype=np.float64) - grad[-1] = 1.0 - - z_axis = grad / np.linalg.norm(grad) - - if D == 2: - x_axis = np.array([z_axis[1], -z_axis[0]], dtype=np.float64) - R = np.column_stack([x_axis, z_axis]) + return evaluate_micro_values_and_gradient( + xyz_to_interpolate, micro_points, micro_weights, anisotropy_matrices, + kernel_range, kernel_type, + )[0] + + +def build_micro_design_matrix(xyz: np.ndarray, centers: np.ndarray, + anisotropy_matrices: np.ndarray, kernel_range: float, + kernel_type: MicroKernelType) -> np.ndarray: + """Rows are evaluation points; column j uses the evaluator's center j metric.""" + if not isinstance(xyz, np.ndarray): + import torch + centers = torch.as_tensor(centers, dtype=xyz.dtype, device=xyz.device) + matrices = torch.as_tensor(anisotropy_matrices, dtype=xyz.dtype, device=xyz.device) + return torch.stack([_kernel_value(torch.linalg.vector_norm((xyz - center) @ matrix.T, dim=1) / kernel_range, + kernel_type) for center, matrix in zip(centers, matrices)], dim=1) + result = np.empty((len(xyz), len(centers)), dtype=np.float64) + for j, (center, matrix) in enumerate(zip(centers, anisotropy_matrices)): + result[:, j] = _kernel_value( + np.linalg.norm((xyz - center) @ matrix.T, axis=1) / kernel_range, kernel_type + ) + return result + + +def evaluate_micro_gradient(xyz: np.ndarray, centers: np.ndarray, weights: np.ndarray, + matrices: np.ndarray, kernel_range: float, + kernel_type: MicroKernelType) -> np.ndarray: + return evaluate_micro_values_and_gradient(xyz, centers, weights, matrices, kernel_range, + kernel_type, compute_gradient=True)[1] + + +def evaluate_micro_values_and_gradient(xyz, centers, weights, matrices, kernel_range=1.0, + kernel_type: MicroKernelType = "exponential", compute_gradient=False): + """Evaluate values and, if requested, gradients in one pass over the centers.""" + if not isinstance(xyz, np.ndarray): + result = _evaluate_micro_torch(xyz, centers, weights, matrices, kernel_range, kernel_type, + compute_gradient=compute_gradient) + return result if compute_gradient else (result, None) + + correction = np.zeros(xyz.shape[0], dtype=np.float64) + gradient = np.zeros_like(xyz, dtype=np.float64) if compute_gradient else None + for center, matrix, weight in zip(centers, matrices, weights): + delta = xyz - center + transformed = delta @ matrix.T + distance = np.linalg.norm(transformed, axis=1) + r = distance / kernel_range + correction += weight * _kernel_value(r, kernel_type) + if compute_gradient: + if kernel_type == "exponential": + derivative = -np.exp(-r) + elif kernel_type == "matern_3_2": + derivative = -3 * r * np.exp(-np.sqrt(3) * r) + elif kernel_type == "matern_5_2": + derivative = -(5 / 3) * r * (1 + np.sqrt(5) * r) * np.exp(-np.sqrt(5) * r) + else: + raise ValueError(f"Unknown micro kernel type: {kernel_type}") + scale = np.divide(derivative, kernel_range * distance, + out=np.zeros_like(distance), where=distance > 0) + gradient += weight * scale[:, None] * (transformed @ matrix) + return correction, gradient + + +def _evaluate_micro_torch(xyz, centers, weights, matrices, kernel_range, kernel_type, + compute_gradient=False): + import torch + + centers = torch.as_tensor(centers, dtype=xyz.dtype, device=xyz.device) + weights = torch.as_tensor(weights, dtype=xyz.dtype, device=xyz.device) + matrices = torch.as_tensor(matrices, dtype=xyz.dtype, device=xyz.device) + correction = xyz.new_zeros(xyz.shape[0]) + gradient = xyz.new_zeros(xyz.shape) if compute_gradient else None + for center, matrix, weight in zip(centers, matrices, weights): + transformed = (xyz - center) @ matrix.T + distance = torch.linalg.vector_norm(transformed, dim=1) + r = distance / kernel_range + if kernel_type == "exponential": + value = torch.exp(-r) + if compute_gradient: + derivative = -value + elif kernel_type == "matern_3_2": + a = np.sqrt(3.) * r + value = (1 + a) * torch.exp(-a) + if compute_gradient: + derivative = -3 * r * torch.exp(-a) + elif kernel_type == "matern_5_2": + a = np.sqrt(5.) * r + value = (1 + a + 5 / 3 * r * r) * torch.exp(-a) + if compute_gradient: + derivative = -(5 / 3) * r * (1 + a) * torch.exp(-a) else: - ref = np.array([0.0, 1.0, 0.0], dtype=np.float64) - if abs(np.dot(z_axis, ref)) > 0.99: - ref = np.array([1.0, 0.0, 0.0], dtype=np.float64) - - x_axis = np.cross(z_axis, ref) - x_axis = x_axis / np.linalg.norm(x_axis) - y_axis = np.cross(z_axis, x_axis) - y_axis = y_axis / np.linalg.norm(y_axis) - - R = np.column_stack([x_axis, y_axis, z_axis]) - - matrices[i] = S @ R.T - - return matrices + raise ValueError(f"Unknown micro kernel type: {kernel_type}") + correction = correction + weight * value + if compute_gradient: + scale = derivative / (kernel_range * distance.clamp_min(torch.finfo(xyz.dtype).tiny)) + gradient = gradient + weight * torch.where(distance > 0, scale, 0)[:, None] * (transformed @ matrix) + return (correction, gradient) if compute_gradient else correction diff --git a/gempy_engine/modules/evaluator/micro_correction.py b/gempy_engine/modules/evaluator/micro_correction.py new file mode 100644 index 00000000..7209de24 --- /dev/null +++ b/gempy_engine/modules/evaluator/micro_correction.py @@ -0,0 +1,100 @@ +"""Fit and apply stack-local authored micro contacts without changing macro solves.""" +from dataclasses import dataclass +import warnings + +import numpy as np + +from gempy_engine.config import AvailableBackends +from gempy_engine.core.backend_tensor import BackendTensor +from gempy_engine.core.data.micro_points import MicroArray +from .micro_anisotropic_evaluator import build_micro_design_matrix, evaluate_micro_values_and_gradient + + +@dataclass(frozen=True) +class MicroCorrection: + points: MicroArray + matrices: MicroArray + weights: MicroArray + kernel_range: float + kernel_type: str + + +def fit_micro_correction(interpolation_input, macro_values, options, surface_sizes): + micro = interpolation_input.micro_points + settings = options.micro_options + if not settings.enabled or micro is None or not len(micro.points): + return None + if (not np.isfinite(settings.kernel_range) or settings.kernel_range <= 0 + or not np.isfinite(settings.nugget) or settings.nugget < 0): + raise ValueError("micro kernel_range must be positive and nuggets nonnegative and finite") + + is_torch = BackendTensor.engine_backend is AvailableBackends.PYTORCH + if is_torch: + import torch + solve_error = torch.linalg.LinAlgError if is_torch else np.linalg.LinAlgError + def tensor(value): + return torch.as_tensor(value, dtype=macro_values.dtype, device=macro_values.device) if is_torch else np.asarray(value) + + macro_points = tensor(interpolation_input.surface_points.sp_coords) + points = tensor(micro.points) + matrices = tensor(micro.anisotropy_matrices) + if is_torch and isinstance(surface_sizes, torch.Tensor): + # Surface counts are integer metadata, not part of the differentiable fit. + sizes = surface_sizes.cpu().tolist() + else: + sizes = np.asarray(surface_sizes, dtype=int).tolist() + if not sizes or sum(sizes) != len(macro_points): + raise ValueError("Micro targets require the stack's surface point counts") + + macro_sp = macro_values[interpolation_input.grid.len_all_grids:interpolation_input.macro_reference_size] + macro_sp = macro_sp[interpolation_input.slice_feature] + contact_values = macro_values[interpolation_input.micro_slice][interpolation_input.micro_indices] + chunks = macro_sp.split(sizes) if is_torch else np.split(macro_sp, np.cumsum(sizes)[:-1]) + means = (torch.stack([chunk.mean() for chunk in chunks]) if is_torch + else np.array([chunk.mean() for chunk in chunks])) + residuals = means[micro.surface_indices] - contact_values + + centers = points + if settings.preserve_macro_points: + # Nearest-center assignment is discrete; gradients flow through selected coordinates and metrics. + nearest = (torch.cdist(macro_points, points).argmin(dim=1) if is_torch else + np.sum((macro_points[:, None] - points[None]) ** 2, axis=2).argmin(axis=1)) + centers = torch.cat((points, macro_points)) if is_torch else np.vstack((points, macro_points)) + matrices = torch.cat((matrices, matrices[nearest])) if is_torch else np.concatenate((matrices, matrices[nearest])) + design = build_micro_design_matrix(centers, centers, matrices, settings.kernel_range, settings.kernel_type) + diagonal = tensor(micro.nuggets) + settings.nugget + if settings.preserve_macro_points: + diagonal = (torch.cat((diagonal, macro_values.new_zeros(len(macro_points)))) if is_torch else + np.concatenate((diagonal, np.zeros(len(macro_points))))) + design = design + (torch.diag(diagonal) if is_torch else np.diag(diagonal)) + rhs = (torch.cat((residuals, macro_values.new_zeros(len(macro_points)))) if is_torch and settings.preserve_macro_points + else np.concatenate((residuals, np.zeros(len(macro_points)))) if settings.preserve_macro_points else residuals) + try: + fitted = torch.linalg.solve(design, rhs) if is_torch else np.linalg.solve(design, rhs) + except solve_error: + fitted = torch.linalg.lstsq(design, rhs).solution if is_torch else np.linalg.lstsq(design, rhs, rcond=None)[0] + warnings.warn("Micro-point system is singular; using least-squares fit", RuntimeWarning, stacklevel=2) + if not (torch.isfinite(fitted).all() if is_torch else np.isfinite(fitted).all()): + raise ValueError("Micro-point fit produced non-finite weights") + return MicroCorrection(centers, matrices, fitted * settings.strength, settings.kernel_range, settings.kernel_type) + + +def apply_micro_correction(fields, xyz, correction: MicroCorrection | None): + if correction is None: + return fields + values, grad = evaluate_micro_values_and_gradient( + xyz, correction.points, correction.weights, correction.matrices, + correction.kernel_range, correction.kernel_type, + compute_gradient=fields.gx_field_everywhere is not None, + ) + if isinstance(fields.scalar_field_everywhere, np.ndarray): + values = values.astype(fields.scalar_field_everywhere.dtype) + fields._scalar_field = fields.scalar_field_everywhere + values + if fields.gx_field_everywhere is not None: + if isinstance(fields.gx_field_everywhere, np.ndarray): + grad = grad.astype(fields.gx_field_everywhere.dtype) + fields._gx_field = fields.gx_field_everywhere + grad[:, 0] + fields._gy_field = fields.gy_field_everywhere + grad[:, 1] + if fields.gz_field_everywhere is not None: + fields._gz_field = fields.gz_field_everywhere + grad[:, 2] + return fields diff --git a/gempy_engine/modules/evaluator/symbolic_evaluator.py b/gempy_engine/modules/evaluator/symbolic_evaluator.py index c3588281..51dcf2c3 100644 --- a/gempy_engine/modules/evaluator/symbolic_evaluator.py +++ b/gempy_engine/modules/evaluator/symbolic_evaluator.py @@ -77,51 +77,9 @@ def symbolic_evaluator(solver_input: SolverInput, weights: np.ndarray, options: else: raise ValueError("Number of dimensions have to be 2 or 3") - scalar_field = _apply_micro_correction(scalar_field, solver_input, options) - return ExportedFields(scalar_field, gx_field, gy_field, gz_field) -def _apply_micro_correction(scalar_field: np.ndarray, solver_input: SolverInput, options: InterpolationOptions) -> np.ndarray: - micro = options.evaluation_options.micro_anisotropic - if not micro.enabled: - return scalar_field - if micro.weights is None or micro.points is None or micro.anisotropy_matrices is None: - return scalar_field - - from gempy_engine.modules.evaluator.micro_anisotropic_evaluator import evaluate_micro_correction - - correction = evaluate_micro_correction( - xyz_to_interpolate=solver_input.xyz_to_interpolate, - micro_points=micro.points, - micro_weights=micro.weights, - anisotropy_matrices=micro.anisotropy_matrices, - kernel_range=micro.kernel_range, - kernel_type=micro.kernel_type, - ) - return scalar_field + correction - - -def _apply_micro_correction_stacked(scalar_field: np.ndarray, eval_input: EvaluatorInput, options: InterpolationOptions) -> np.ndarray: - micro = options.evaluation_options.micro_anisotropic - if not micro.enabled: - return scalar_field - if micro.weights is None or micro.points is None or micro.anisotropy_matrices is None: - return scalar_field - - from gempy_engine.modules.evaluator.micro_anisotropic_evaluator import evaluate_micro_correction - - correction = evaluate_micro_correction( - xyz_to_interpolate=eval_input.xyz_to_interpolate, - micro_points=micro.points, - micro_weights=micro.weights, - anisotropy_matrices=micro.anisotropy_matrices, - kernel_range=micro.kernel_range, - kernel_type=micro.kernel_type, - ) - return scalar_field + correction - - def _build_block_sparse_ranges(M_sizes: list[int], N_sizes: list[int]): """Build PyKeOps block-sparse ranges tuple for block-diagonal evaluation.""" keep_i = np.cumsum([0] + N_sizes) @@ -164,7 +122,7 @@ def _validate_stacked_dimensions(eval_kernel, weights, M_sizes: list[int], N_siz def symbolic_evaluator_optimized_stacked( eval_inputs: list[EvaluatorInput], weights_list: list[np.ndarray], - options_list: list[InterpolationOptions] + options_list: list[InterpolationOptions], ) -> list[ExportedFields]: """Evaluate multiple fields in a single PyKeOps call using block-sparse ranges. @@ -334,9 +292,8 @@ def _run_prep(args): if gy_field is not None: gy_field = BackendTensor.t.to_numpy(gy_field) if gz_field is not None: gz_field = BackendTensor.t.to_numpy(gz_field) - s_field = _apply_micro_correction_stacked(s_field, eval_inputs[idx], options_list[idx]) - - results.append(ExportedFields(s_field, gx_field, gy_field, gz_field)) + fields = ExportedFields(s_field, gx_field, gy_field, gz_field) + results.append(fields) return results diff --git a/tests/test_common/test_api/test_concurrent_options.py b/tests/test_common/test_api/test_concurrent_options.py index 61a31eab..ae0ff6ee 100644 --- a/tests/test_common/test_api/test_concurrent_options.py +++ b/tests/test_common/test_api/test_concurrent_options.py @@ -1,5 +1,6 @@ from concurrent.futures import ThreadPoolExecutor from threading import Barrier +from types import SimpleNamespace import pytest @@ -24,7 +25,7 @@ def test_computations_do_not_share_volatile_option_state(monkeypatch): monkeypatch.setattr(model_api, "_check_input_validity", lambda *_: None) def observe_options(interpolation_input, options, data_descriptor): - level = interpolation_input + level = interpolation_input.level options.temp_interpolation_values.current_octree_level = level barrier.wait() return options.temp_interpolation_values.current_octree_level @@ -39,7 +40,9 @@ def __init__(self, octrees_output, **_): with ThreadPoolExecutor(max_workers=2) as executor: futures = [ - executor.submit(model_api.compute_model, level, options, None) + executor.submit(model_api.compute_model, + SimpleNamespace(level=level, micro_points=None), options, + SimpleNamespace(stack_structure=None)) for level in (1, 3) ] diff --git a/tests/test_common/test_api/test_micro_points_integration.py b/tests/test_common/test_api/test_micro_points_integration.py new file mode 100644 index 00000000..0aaf0bae --- /dev/null +++ b/tests/test_common/test_api/test_micro_points_integration.py @@ -0,0 +1,495 @@ +import numpy as np +import pytest + +from gempy_engine.API.model.model_api import compute_model +from gempy_engine.core.data import InterpolationOptions, Orientations, SurfacePoints, TensorsStructure +from gempy_engine.core.data.engine_grid import EngineGrid, RegularGrid +from gempy_engine.core.data.input_data_descriptor import InputDataDescriptor +from gempy_engine.core.data.interpolation_input import InterpolationInput +from gempy_engine.core.data.micro_points import MicroPoints +from gempy_engine.core.data.stack_relation_type import StackRelationType +from gempy_engine.core.data.stacks_structure import StacksStructure +from gempy_engine.core.backend_tensor import BackendTensor +from gempy_engine.config import AvailableBackends +from gempy_engine.core.data.kernel_classes.faults import FaultsData +from gempy_engine.core.data.finite_fault import FiniteFault +from gempy_engine.core.data.interpolation_functions import CustomInterpolationFunctions +from gempy_engine.modules.evaluator.micro_anisotropic_evaluator import ( + build_micro_design_matrix, evaluate_micro_correction, evaluate_micro_gradient, +) + + +@pytest.mark.parametrize("flat", [False, True]) +@pytest.mark.parametrize("gradient", [False, True]) +def test_authored_micro_contacts_are_stack_local(monkeypatch, flat, gradient): + if flat: + BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=True) + monkeypatch.setenv("GEMPY_FLAT_STACKS", "True") + try: + _assert_authored_micro_contacts_are_stack_local(monkeypatch, gradient) + finally: + if flat: + BackendTensor._change_backend(AvailableBackends.numpy) + + +def _assert_authored_micro_contacts_are_stack_local(monkeypatch, gradient): + sp = np.array([[0.2, 0.2, 0.4], [0.8, 0.2, 0.4], [0.2, 0.8, 0.4], + [0.8, 0.8, 0.4]]) + points = np.vstack((sp, sp + [0, 0, 0.2])) + orientations = Orientations(dip_positions=np.array([[0.5, 0.5, 0.5], [0.5, 0.5, 0.7]]), + dip_gradients=np.array([[0, 0, 1], [0, 0, 1]])) + grid = EngineGrid.from_regular_grid(RegularGrid( + orthogonal_extent=[0, 1, 0, 1, 0, 1], regular_grid_shape=[3, 3, 3])) + contact = np.array([[0.5, 0.5, 0.52]]) + micro = MicroPoints(contact, np.array([np.diag([2., 2., 2.])]), + np.array([0.0]), np.array([0])) + ii = InterpolationInput(SurfacePoints(points), orientations, grid, micro_points=micro) + options = InterpolationOptions.from_args(range=3., c_o=1., uni_degree=0, + mesh_extraction=False) + options.evaluation_options.compute_scalar_gradient = gradient + descriptor = InputDataDescriptor(TensorsStructure(np.array([4, 4])), StacksStructure( + number_of_points_per_stack=np.array([4, 4]), + number_of_orientations_per_stack=np.array([1, 1]), + number_of_surfaces_per_stack=np.array([1, 1]), + masking_descriptor=[StackRelationType.ERODE, StackRelationType.ERODE])) + + def fields(): + result = compute_model(ii, options, descriptor) + if BackendTensor.engine_backend is AvailableBackends.PYTORCH: + for output in result.octrees_output[0].outputs: + field = output.scalar_fields.exported_fields.scalar_field_everywhere + assert field.device == ii.surface_points.sp_coords.device + assert field.dtype == ii.surface_points.sp_coords.dtype + return [BackendTensor.t.to_numpy(output.scalar_fields.exported_fields.scalar_field_everywhere).copy() + for output in result.octrees_output[0].outputs] + + with pytest.warns(UserWarning, match="Stack 0 contains micro points.*ignored") as recorded: + baseline = fields() + assert len(recorded) == 1 + options.micro_options.enabled = True + monkeypatch.setenv("GEMPY_FLAT_STACKS", "True") + corrected = fields() + if BackendTensor.use_pykeops: + monkeypatch.setenv("GEMPY_FLAT_STACKS", "False") + sequential = fields() + for flat_field, sequential_field in zip(corrected, sequential): + np.testing.assert_allclose(flat_field, sequential_field, atol=1e-5, rtol=1e-5) + assert len(corrected[0]) == len(baseline[0]) + 1 + assert np.max(np.abs(corrected[0][:-1] - baseline[0])) > 1e-6 + np.testing.assert_allclose(corrected[1][:-1], baseline[1], atol=1e-8) + grid_size = grid.len_all_grids + np.testing.assert_allclose(corrected[0][grid_size:grid_size + 4], + baseline[0][grid_size:grid_size + 4], atol=1e-6) + assert corrected[0].shape == corrected[1].shape + assert not hasattr(options.micro_options, "weights") + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_torch_authored_fit_preserves_autograd(device): + torch = pytest.importorskip("torch") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + BackendTensor._change_backend(AvailableBackends.PYTORCH, use_gpu=device == "cuda", grads=True) + try: + sp = torch.tensor([[0.2, 0.2, 0.4], [0.8, 0.2, 0.4], + [0.2, 0.8, 0.4], [0.8, 0.8, 0.4]], dtype=torch.float32, requires_grad=True) + contact = torch.tensor([[0.5, 0.5, 0.52]], dtype=torch.float32, requires_grad=True) + metric = torch.diag(torch.tensor([2., 2., 2.], dtype=torch.float32)).unsqueeze(0).requires_grad_() + nugget = torch.tensor([0.01], dtype=torch.float32, requires_grad=True) + ii = InterpolationInput(SurfacePoints(sp), Orientations( + dip_positions=np.array([[0.5, 0.5, 0.5]]), dip_gradients=np.array([[0., 0., 1.]])), + EngineGrid.from_regular_grid(RegularGrid( + orthogonal_extent=[0, 1, 0, 1, 0, 1], regular_grid_shape=[3, 3, 3])), + micro_points=MicroPoints(contact, metric, nugget, np.array([0]))) + options = InterpolationOptions.from_args(range=3., c_o=1., uni_degree=0, mesh_extraction=False) + options.micro_options.enabled = True + options.evaluation_options.compute_scalar_gradient = True + descriptor = InputDataDescriptor(TensorsStructure(np.array([4])), StacksStructure( + np.array([4]), np.array([1]), np.array([1]), [StackRelationType.ERODE])) + result = compute_model(ii, options, descriptor) + field = result.octrees_output[0].outputs[0].scalar_fields.exported_fields.scalar_field_everywhere + assert field.requires_grad + grads = torch.autograd.grad(field[:ii.grid.len_all_grids].sum(), (sp, contact, metric, nugget)) + for grad in grads: + assert torch.isfinite(grad).all() and grad.abs().max() > 1e-9 + assert ii.micro_points.points is contact + assert not hasattr(options.micro_options, "weights") + finally: + torch.set_default_device("cpu") + BackendTensor._change_backend(AvailableBackends.numpy) + + +def test_micro_points_reject_invalid_surface_index(): + with pytest.raises(ValueError, match="surface_indices"): + MicroPoints(np.zeros((1, 3)), np.eye(3)[None], np.zeros(1), np.array([0.5])) + + +def test_unknown_global_surface_index_with_stack_overrides(): + points = np.array([[0.2, 0.2, 0.4], [0.8, 0.2, 0.4], + [0.2, 0.8, 0.4], [0.8, 0.8, 0.4]]) + ii = InterpolationInput( + SurfacePoints(points), + Orientations(np.array([[0.5, 0.5, 0.5]]), np.array([[0., 0., 1.]])), + EngineGrid.from_regular_grid(RegularGrid( + orthogonal_extent=[0, 1, 0, 1, 0, 1], regular_grid_shape=[2, 2, 2])), + micro_points=MicroPoints(np.array([[0.5, 0.5, 0.52]]), np.eye(3)[None], + np.zeros(1), np.array([1])), + ) + options = InterpolationOptions.from_args(range=3., c_o=1., uni_degree=0, mesh_extraction=False) + overrides = [InterpolationOptions.from_args(range=3., c_o=1., uni_degree=0, mesh_extraction=False)] + descriptor = InputDataDescriptor(TensorsStructure(np.array([4])), StacksStructure( + np.array([4]), np.array([1]), np.array([1]), [StackRelationType.ERODE], + interpolation_options_per_stack=overrides)) + with pytest.raises(ValueError, match="unknown global surface index"): + compute_model(ii, options, descriptor) + + +def test_micro_subset_uses_global_surface_indices(): + grid = EngineGrid.from_regular_grid(RegularGrid( + orthogonal_extent=[0, 1, 0, 1, 0, 1], regular_grid_shape=[2, 2, 2])) + contacts = MicroPoints(np.array([[0., 0., 0.], [1., 1., 1.]]), + np.array([np.eye(3), np.eye(3)]), np.zeros(2), np.array([0, 1])) + ii = InterpolationInput(SurfacePoints(np.zeros((4, 3))), + Orientations(np.zeros((2, 3)), np.ones((2, 3))), grid, + micro_points=contacts) + stacks = StacksStructure(np.array([2, 2]), np.array([1, 1]), np.array([1, 1]), + [StackRelationType.ERODE, StackRelationType.ERODE]) + stacks.stack_number = 1 + subset = InterpolationInput.from_interpolation_input_subset(ii, stacks) + np.testing.assert_array_equal(subset.micro_points.points, contacts.points[1:]) + np.testing.assert_array_equal(subset.micro_points.surface_indices, [0]) + ii._all_micro_points = contacts + subset = InterpolationInput.from_interpolation_input_subset(ii, stacks) + np.testing.assert_array_equal(subset.micro_indices, [1]) + assert subset.micro_slice == slice(grid.len_all_grids + 4, grid.len_all_grids + 6) + + +def test_micro_design_and_gradient_match_evaluation(): + centers = np.array([[0., 0., 0.], [1., 0., 0.]]) + matrices = np.array([np.diag([1., 2., 3.]), np.diag([3., 1., 2.])]) + weights = np.array([0.3, -0.7]) + xyz = np.array([[0.4, 0.2, 0.1], [0.8, 0.3, 0.2]]) + design = build_micro_design_matrix(centers, centers, matrices, 0.6, "matern_5_2") + for j in range(len(centers)): + np.testing.assert_allclose(design[:, j], evaluate_micro_correction( + centers, centers, np.eye(len(centers))[j], matrices, 0.6, "matern_5_2")) + grad = evaluate_micro_gradient(xyz, centers, weights, matrices, 0.6, "matern_5_2") + for axis in range(3): + shift = np.eye(3)[axis] * 1e-6 + difference = (evaluate_micro_correction(xyz + shift, centers, weights, matrices, 0.6, "matern_5_2") + - evaluate_micro_correction(xyz - shift, centers, weights, matrices, 0.6, "matern_5_2")) / 2e-6 + np.testing.assert_allclose(grad[:, axis], difference, atol=1e-8) + + +@pytest.mark.parametrize("kernel", ["exponential", "matern_3_2", "matern_5_2"]) +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_torch_micro_field_matches_numpy_and_autograd(kernel, dtype, device): + torch = pytest.importorskip("torch") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + centers = np.array([[0., 0., 0.], [1., 0., 0.]]) + matrices = np.array([np.diag([1., 2., 3.]), np.diag([3., 1., 2.])]) + coords = np.array([[0.4, 0.2, 0.1], [0.8, 0.3, 0.2]]) + with torch.enable_grad(): + xyz = torch.tensor(coords, dtype=getattr(torch, dtype), device=device, requires_grad=True) + weights = torch.tensor([0.3, -0.7], dtype=xyz.dtype, device=device, requires_grad=True) + correction = evaluate_micro_correction(xyz, centers, weights, matrices, 0.6, kernel) + gradient = evaluate_micro_gradient(xyz, centers, weights, matrices, 0.6, kernel) + xyz_grad, weight_grad = torch.autograd.grad(correction.sum(), (xyz, weights)) + assert correction.dtype == gradient.dtype == xyz.dtype + assert correction.device == gradient.device == xyz.device + np.testing.assert_allclose(correction.detach().cpu(), evaluate_micro_correction( + coords, centers, weights.detach().cpu().numpy(), matrices, 0.6, kernel), rtol=1e-5, atol=1e-6) + np.testing.assert_allclose(gradient.detach().cpu(), xyz_grad.detach().cpu(), rtol=1e-5, atol=1e-6) + np.testing.assert_allclose(weight_grad.detach().cpu(), build_micro_design_matrix( + coords, centers, matrices, 0.6, kernel).sum(axis=0), rtol=1e-5, atol=1e-6) + + +def _fault_micro_model(finite=False, chain=False): + fault_sp = np.array([[0.5, 0.2, 0.2], [0.5, 0.8, 0.2], + [0.5, 0.2, 0.8], [0.5, 0.8, 0.8]]) + strata_sp = np.array([[0.2, 0.2, 0.4], [0.8, 0.2, 0.4], + [0.2, 0.8, 0.4], [0.8, 0.8, 0.4]]) + points = [fault_sp] + if chain: + points.append(fault_sp + [0.12, 0, 0]) + points.append(strata_sp) + count = len(points) + relations = np.zeros((count, count), dtype=bool) + relations[0, -1] = True + if chain: + relations[0, 1] = relations[1, 2] = True + contacts = np.array([[0.56, 0.5, 0.54], [0.5, 0.5, 0.5], fault_sp[0]]) + micro = MicroPoints(contacts, np.repeat(np.eye(3)[None] * 3, 3, axis=0), + np.zeros(3), np.array([count - 1] * 3)) + grid = EngineGrid.from_regular_grid(RegularGrid( + orthogonal_extent=[0, 1, 0, 1, 0, 1], regular_grid_shape=[3, 3, 3])) + ii = InterpolationInput(SurfacePoints(np.vstack(points)), Orientations( + np.array([[0.5, 0.5, 0.5]] * (count - 1) + [[0.5, 0.5, 0.4]]), + np.array([[1., 0., 0.]] * (count - 1) + [[0., 0., 1.]])), + grid, micro_points=micro) + faults = [None] * count + if finite: + faults[0] = FaultsData.from_user_input(None, FiniteFault( + center=(0.5, 0.5, 0.5), strike_radius=0.8, dip_radius=0.8)) + descriptor = InputDataDescriptor(TensorsStructure(np.array([4] * count)), StacksStructure( + np.array([4] * count), np.array([1] * count), np.array([1] * count), + [StackRelationType.FAULT] * (count - 1) + [StackRelationType.ERODE], + faults_relations=relations, faults_input_data=faults)) + options = InterpolationOptions.from_args(range=3., c_o=1., uni_degree=0, mesh_extraction=False) + options.micro_options.enabled = True + options.evaluation_options.number_octree_levels = 1 + return ii, descriptor, options + + +@pytest.mark.parametrize('finite,chain', [(False, False), (False, True), (True, False)]) +def test_fault_micro_contacts_flow_through_dependencies(monkeypatch, finite, chain): + ii, descriptor, options = _fault_micro_model(finite, chain) + from gempy_engine.API.interp_single._aux_faults_ops import _modify_faults_values_output + + normalized = [] + original = _modify_faults_values_output + + def capture(fault_input, output, xyz_to_interpolate): + values = original(fault_input, output, xyz_to_interpolate) + normalized.append(np.asarray(values).copy()) + return values + + monkeypatch.setattr('gempy_engine.API.interp_single._multi_scalar_field_manager._modify_faults_values_output', capture) + monkeypatch.setattr('gempy_engine.API.interp_single._stack_ops._modify_faults_values_output', capture) + authored_points = ii.micro_points.points.copy() + surface_points = ii.surface_points.sp_coords.copy() + def run(flat, enabled=True): + options.micro_options.enabled = enabled + monkeypatch.setenv('GEMPY_FLAT_STACKS', str(flat)) + BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=flat) + result = compute_model(ii, options, descriptor) + outputs = result.octrees_output[0].outputs + return [np.asarray(o.scalar_fields.exported_fields.scalar_field_everywhere).copy() for o in outputs], outputs + + try: + serial, _ = run(False) + assert len(normalized) == len(serial) - 1 + assert all(v.shape == (1, len(serial[0])) and np.isfinite(v).all() for v in normalized) + assert np.any(normalized[0][0, -3:] != normalized[0][0, -4]) + normalized.clear() + flat, _ = run(True) + for left, right in zip(serial, flat): + np.testing.assert_allclose(left, right, rtol=1e-4, atol=1e-4) + baseline, baseline_outputs = run(False, False) + np.testing.assert_allclose(serial[0][:-3], baseline[0], atol=1e-8) + assert np.max(np.abs(serial[-1][:-3] - baseline[-1])) > 1e-5 + grid_index = np.flatnonzero(np.all(np.isclose(ii.grid.values, [0.5, 0.5, 0.5]), axis=1))[0] + target = np.asarray(baseline_outputs[-1].scalar_fields.scalar_field_at_sp).item() + assert abs(serial[-1][grid_index] - target) < abs(baseline[-1][grid_index] - target) + np.testing.assert_array_equal(ii.micro_points.points, authored_points) + np.testing.assert_array_equal(ii.surface_points.sp_coords, surface_points) + finally: + BackendTensor._change_backend(AvailableBackends.numpy) + + +def test_fault_micro_query_preserves_torch_gradients(): + torch = pytest.importorskip('torch') + ii, descriptor, options = _fault_micro_model() + sp = torch.tensor(ii.surface_points.sp_coords, dtype=torch.float32, requires_grad=True) + contacts = torch.tensor(ii.micro_points.points, dtype=torch.float32, requires_grad=True) + ii.surface_points = SurfacePoints(sp) + ii.micro_points = MicroPoints(contacts, ii.micro_points.anisotropy_matrices, + ii.micro_points.nuggets, ii.micro_points.surface_indices) + BackendTensor._change_backend(AvailableBackends.PYTORCH, grads=True) + try: + result = compute_model(ii, options, descriptor) + field = result.octrees_output[0].outputs[-1].exported_fields.scalar_field + grads = torch.autograd.grad(field.sum(), (sp, contacts)) + assert all(torch.isfinite(grad).all() and grad.abs().max() > 1e-8 for grad in grads) + finally: + BackendTensor._change_backend(AvailableBackends.numpy) + + +@pytest.mark.parametrize('finite', [False, True]) +@pytest.mark.parametrize('flat', [False, True]) +def test_fault_reference_prefix_unaffected_by_remote_contact(monkeypatch, finite, flat): + ii, descriptor, options = _fault_micro_model(finite=finite) + ii.micro_points.points[0] = [50., -50., 50.] + from gempy_engine.API.interp_single._aux_faults_ops import _modify_faults_values_output + rows = [] + + def capture(fault_input, output, xyz_to_interpolate): + values = _modify_faults_values_output(fault_input, output, xyz_to_interpolate) + rows.append(np.asarray(values).copy()) + return values + + module = ('_stack_ops' if flat else '_multi_scalar_field_manager') + monkeypatch.setattr(f'gempy_engine.API.interp_single.{module}._modify_faults_values_output', capture) + monkeypatch.setenv('GEMPY_FLAT_STACKS', str(flat)) + BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=flat) + try: + compute_model(ii, options, descriptor) + enabled = rows.pop() + options.micro_options.enabled = False + with pytest.warns(UserWarning, match='ignored'): + compute_model(ii, options, descriptor) + disabled = rows.pop() + np.testing.assert_allclose(enabled[:, :-3], disabled, rtol=1e-6, atol=1e-6) + assert enabled.shape[1] == disabled.shape[1] + 3 + finally: + BackendTensor._change_backend(AvailableBackends.numpy) + + +def test_micro_fit_uses_existing_macro_evaluation(monkeypatch): + ii, descriptor, options = _fault_micro_model(chain=True) + from gempy_engine.API.interp_single import _interp_single_feature + original = _interp_single_feature._evaluate_sys_eq + calls = [] + + def capture(*args, **kwargs): + calls.append(args[0].xyz_to_interpolate.shape[0]) + return original(*args, **kwargs) + + monkeypatch.setenv('GEMPY_FLAT_STACKS', 'False') + monkeypatch.setattr(_interp_single_feature, '_evaluate_sys_eq', capture) + compute_model(ii, options, descriptor) + assert calls == [ii.grid.len_all_grids + ii.surface_points.n_points + len(ii.micro_points.points)] * 3 + + +@pytest.mark.parametrize('flat', [False, True]) +def test_external_upstream_fault_evaluates_shared_suffix(monkeypatch, flat): + ii, descriptor, options = _fault_micro_model() + descriptor.stack_structure.interp_functions_per_stack = [CustomInterpolationFunctions( + scalar_field_at_surface_points=np.array([0.5]), + implicit_function=lambda xyz: xyz[:, 0], + ), None] + monkeypatch.setenv('GEMPY_FLAT_STACKS', str(flat)) + BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=flat) + try: + result = compute_model(ii, options, descriptor) + outputs = result.octrees_output[0].outputs + expected = ii.grid.len_all_grids + ii.surface_points.n_points + len(ii.micro_points.points) + assert all(len(o.scalar_fields.exported_fields.scalar_field_everywhere) == expected for o in outputs) + assert np.isfinite(outputs[-1].scalar_fields.exported_fields.scalar_field_everywhere).all() + finally: + BackendTensor._change_backend(AvailableBackends.numpy) + + +def test_flat_faults_use_each_faults_segmentation_function(monkeypatch): + ii, descriptor, options = _fault_micro_model(chain=True) + relations = descriptor.stack_structure.faults_relations + relations[0, 1] = False + relations[1, 2] = False + descriptor.stack_structure.segmentation_functions_per_stack = [ + lambda xyz: 0.5, lambda xyz: 30., None, + ] + from gempy_engine.API.interp_single._aux_faults_ops import _modify_faults_values_output + published = [] + def capture(fault_input, output, xyz_to_interpolate): + values = _modify_faults_values_output(fault_input, output, xyz_to_interpolate) + published.append(np.asarray(values).copy()) + return values + monkeypatch.setattr('gempy_engine.API.interp_single._multi_scalar_field_manager._modify_faults_values_output', capture) + monkeypatch.setattr('gempy_engine.API.interp_single._stack_ops._modify_faults_values_output', capture) + try: + BackendTensor._change_backend(AvailableBackends.numpy) + monkeypatch.setenv('GEMPY_FLAT_STACKS', 'False') + serial = compute_model(ii, options, descriptor) + serial_queries = published.copy() + published.clear() + BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=True) + monkeypatch.setenv('GEMPY_FLAT_STACKS', 'True') + flat = compute_model(ii, options, descriptor) + assert len(serial_queries) == len(published) == 2 + for expected, actual in zip(serial_queries, published): + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose( + flat.octrees_output[0].outputs[-1].scalar_fields.exported_fields.scalar_field_everywhere, + serial.octrees_output[0].outputs[-1].scalar_fields.exported_fields.scalar_field_everywhere, + rtol=1e-5, atol=1e-5, + ) + finally: + BackendTensor._change_backend(AvailableBackends.numpy) + + +@pytest.mark.parametrize('flat', [False, True]) +@pytest.mark.parametrize('chain', [False, True]) +def test_fault_surface_micro_rejected_even_with_upstream_faults(monkeypatch, flat, chain): + ii, descriptor, options = _fault_micro_model(chain=chain) + ii.micro_points.surface_indices[0] = 1 if chain else 0 + monkeypatch.setenv('GEMPY_FLAT_STACKS', str(flat)) + BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=flat) + try: + with pytest.raises(ValueError, match='fault stack.*micro points on fault surfaces'): + compute_model(ii, options, descriptor) + finally: + BackendTensor._change_backend(AvailableBackends.numpy) + + +@pytest.mark.parametrize('flat', [False, True]) +@pytest.mark.parametrize('root_enabled', [False, True]) +def test_disabled_fault_surface_micro_is_ignored(monkeypatch, flat, root_enabled): + ii, descriptor, options = _fault_micro_model() + ii.micro_points.surface_indices[0] = 0 + overrides = [options.model_copy(deep=True), options.model_copy(deep=True)] + overrides[0].micro_options.enabled = False + descriptor.stack_structure.interpolation_options_per_stack = overrides + options.micro_options.enabled = root_enabled + monkeypatch.setenv('GEMPY_FLAT_STACKS', str(flat)) + BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=flat) + try: + with pytest.warns(UserWarning, match='Stack 0 contains micro points.*ignored'): + result = compute_model(ii, options, descriptor) + assert len(result.octrees_output[0].outputs) == 2 + expected_size = ii.grid.len_all_grids + ii.surface_points.n_points + len(ii.micro_points.points) + assert all(len(o.scalar_fields.exported_fields.scalar_field_everywhere) == expected_size + for o in result.octrees_output[0].outputs) + finally: + BackendTensor._change_backend(AvailableBackends.numpy) + + +@pytest.mark.parametrize('flat', [False, True]) +def test_external_micro_contacts_rejected_before_evaluation(monkeypatch, flat): + ii, descriptor, options = _fault_micro_model() + descriptor.stack_structure.interp_functions_per_stack = [None, CustomInterpolationFunctions( + scalar_field_at_surface_points=np.array([0.4]), + implicit_function=lambda xyz: xyz[:, 2], + )] + monkeypatch.setenv('GEMPY_FLAT_STACKS', str(flat)) + + def unexpected_evaluation(*args, **kwargs): + pytest.fail('Unsupported micro contacts must be rejected before interpolation') + + monkeypatch.setattr('gempy_engine.API.model.model_api.interpolate_n_octree_levels', unexpected_evaluation) + with pytest.raises(NotImplementedError, match='external-function stacks'): + compute_model(ii, options, descriptor) + + +@pytest.mark.parametrize('flat', [False, True]) +@pytest.mark.parametrize('deduplicate', [False, True]) +def test_micro_contacts_flow_through_octree_and_mesh(monkeypatch, flat, deduplicate): + ii, descriptor, options = _fault_micro_model() + ii.set_temp_grid(EngineGrid.from_regular_grid(RegularGrid( + orthogonal_extent=[0, 1, 0, 1, 0, 1], regular_grid_shape=[2, 2, 2]))) + options.evaluation_options.number_octree_levels = 2 + options.evaluation_options.number_octree_levels_surface = 2 + options.evaluation_options.mesh_extraction = True + options.evaluation_options.compute_scalar_gradient = True + options.evaluation_options.deduplicate_octree_corners = deduplicate + monkeypatch.setenv('GEMPY_FLAT_STACKS', str(flat)) + BackendTensor._change_backend(AvailableBackends.numpy, use_pykeops=flat) + try: + result = compute_model(ii, options, descriptor) + assert len(result.octrees_output) == 2 + for level in result.octrees_output: + for output in level.outputs: + fields = output.exported_fields + assert len(fields.scalar_field) == output.grid.len_all_grids + assert len(fields.scalar_field_everywhere) == ( + output.grid.len_all_grids + ii.surface_points.n_points + len(ii.micro_points.points)) + assert np.isfinite(fields.scalar_field_everywhere).all() + assert result.dc_meshes + for mesh in result.dc_meshes: + assert mesh is not None + assert len(mesh.vertices) > 0 + assert len(mesh.edges) > 0 + assert np.isfinite(mesh.vertices).all() + finally: + BackendTensor._change_backend(AvailableBackends.numpy) diff --git a/tests/test_common/test_modules/test_evaluator/micro_reference.py b/tests/test_common/test_modules/test_evaluator/micro_reference.py new file mode 100644 index 00000000..232c2ea7 --- /dev/null +++ b/tests/test_common/test_modules/test_evaluator/micro_reference.py @@ -0,0 +1,94 @@ +"""Dense reference micro solve used by evaluator tests, not the production fit.""" + +import numpy as np + +from gempy_engine.modules.evaluator.micro_anisotropic_evaluator import _kernel_value, MicroKernelType + + +def build_micro_covariance( + micro_points: np.ndarray, + anisotropy_matrices: np.ndarray, + kernel_range: float = 1.0, + kernel_type: MicroKernelType = "exponential", + nugget: float = 0.0, +) -> np.ndarray: + """Build the symmetric NxN covariance matrix for the reference micro solve. + + Distances use the symmetric metric M_ij = (A_i^T A_i + A_j^T A_j) / 2. + """ + N = micro_points.shape[0] + K = np.zeros((N, N), dtype=np.float64) + ATA = np.einsum('nki,nkj->nij', anisotropy_matrices, anisotropy_matrices) + + for i in range(N): + for j in range(i, N): + M_ij = 0.5 * (ATA[i] + ATA[j]) + diff = micro_points[i] - micro_points[j] + dist_sq = diff @ M_ij @ diff + dist = np.sqrt(max(dist_sq, 0.0)) + r = dist / kernel_range + val = float(_kernel_value(np.array(r), kernel_type)) + K[i, j] = val + K[j, i] = val + + if nugget > 0: + np.fill_diagonal(K, K.diagonal() + nugget) + return K + + +def solve_micro_weights( + micro_points: np.ndarray, + residuals: np.ndarray, + anisotropy_matrices: np.ndarray, + kernel_range: float = 1.0, + kernel_type: MicroKernelType = "exponential", + nugget: float = 0.0, +) -> np.ndarray: + """Solve the dense reference covariance system.""" + K = build_micro_covariance(micro_points, anisotropy_matrices, kernel_range, kernel_type, nugget) + weights = np.linalg.solve(K, residuals) + return weights + + +def compute_anisotropy_matrices_from_gradients( + micro_points: np.ndarray, + gradients: np.ndarray, + r_vertical: float = 1.0, + r_lateral: float = 10.0, +) -> np.ndarray: + """Build per-point anisotropy matrices from macro gradient directions (2D/3D).""" + N, D = micro_points.shape + assert gradients.shape == (N, D), f"gradients shape {gradients.shape} != (N, D) {(N, D)}" + + if D == 2: + lateral_scales = np.array([1.0 / r_lateral], dtype=np.float64) + scales = np.concatenate([lateral_scales, [1.0 / r_vertical]]) + S = np.diag(scales) + else: + S = np.diag(np.array([1.0 / r_lateral, 1.0 / r_lateral, 1.0 / r_vertical])) + + matrices = np.zeros((N, D, D), dtype=np.float64) + for i in range(N): + grad = gradients[i].astype(np.float64) + grad_norm = np.linalg.norm(grad) + if grad_norm < 1e-10: + grad = np.zeros(D, dtype=np.float64) + grad[-1] = 1.0 + + z_axis = grad / np.linalg.norm(grad) + if D == 2: + x_axis = np.array([z_axis[1], -z_axis[0]], dtype=np.float64) + R = np.column_stack([x_axis, z_axis]) + else: + ref = np.array([0.0, 1.0, 0.0], dtype=np.float64) + if abs(np.dot(z_axis, ref)) > 0.99: + ref = np.array([1.0, 0.0, 0.0], dtype=np.float64) + + x_axis = np.cross(z_axis, ref) + x_axis = x_axis / np.linalg.norm(x_axis) + y_axis = np.cross(z_axis, x_axis) + y_axis = y_axis / np.linalg.norm(y_axis) + R = np.column_stack([x_axis, y_axis, z_axis]) + + matrices[i] = S @ R.T + return matrices diff --git a/tests/test_common/test_modules/test_evaluator/test_micro_anisotropic_evaluator.py b/tests/test_common/test_modules/test_evaluator/test_micro_anisotropic_evaluator.py index 7b5e4e41..d694d76c 100644 --- a/tests/test_common/test_modules/test_evaluator/test_micro_anisotropic_evaluator.py +++ b/tests/test_common/test_modules/test_evaluator/test_micro_anisotropic_evaluator.py @@ -1,7 +1,12 @@ import numpy as np +import pytest from gempy_engine.modules.evaluator.micro_anisotropic_evaluator import ( evaluate_micro_correction, + evaluate_micro_gradient, + evaluate_micro_values_and_gradient, +) +from .micro_reference import ( build_micro_covariance, solve_micro_weights, compute_anisotropy_matrices_from_gradients, @@ -12,6 +17,49 @@ def _make_identity_anisotropy(N: int) -> np.ndarray: return np.tile(np.eye(3, dtype=np.float64), (N, 1, 1)) +@pytest.mark.parametrize("kernel_type", ["exponential", "matern_3_2", "matern_5_2"]) +def test_joint_values_and_gradient_numpy(kernel_type): + xyz = np.array([[0., 0., 0.], [.7, -.4, .2], [1., 2., 3.]]) + centers = np.array([[0., 0., 0.], [1., 1., 1.]]) + weights = np.array([1.2, -.4]) + matrices = np.array([np.diag([2., 1., .5]), np.diag([.5, 2., 1.])]) + values, gradient = evaluate_micro_values_and_gradient( + xyz, centers, weights, matrices, 1.3, kernel_type, compute_gradient=True, + ) + only_values, no_gradient = evaluate_micro_values_and_gradient(xyz, centers, weights, matrices, 1.3, kernel_type) + np.testing.assert_allclose(values, only_values) + np.testing.assert_allclose(values, evaluate_micro_correction(xyz, centers, weights, matrices, 1.3, kernel_type)) + np.testing.assert_allclose(gradient, evaluate_micro_gradient(xyz, centers, weights, matrices, 1.3, kernel_type)) + assert no_gradient is None + for axis in range(3): + shifted = xyz.copy() + shifted[:, axis] += 1e-6 + finite_difference = (evaluate_micro_correction(shifted, centers, weights, matrices, 1.3, kernel_type) + - values) / 1e-6 + np.testing.assert_allclose(gradient[1:, axis], finite_difference[1:], atol=2e-6) + + +@pytest.mark.parametrize("kernel_type", ["exponential", "matern_3_2", "matern_5_2"]) +def test_joint_values_and_gradient_torch(kernel_type): + torch = pytest.importorskip("torch") + xyz = torch.tensor([[0., 0., 0.], [.7, -.4, .2]], dtype=torch.float64, requires_grad=True) + centers = torch.tensor([[0., 0., 0.]], dtype=torch.float64) + weights = torch.tensor([1.2], dtype=torch.float64, requires_grad=True) + matrices = torch.diag(torch.tensor([2., 1., .5], dtype=torch.float64))[None] + values, gradient = evaluate_micro_values_and_gradient( + xyz, centers, weights, matrices, 1.3, kernel_type, compute_gradient=True, + ) + only_values, no_gradient = evaluate_micro_values_and_gradient(xyz, centers, weights, matrices, 1.3, kernel_type) + torch.testing.assert_close(values, only_values) + torch.testing.assert_close(values, evaluate_micro_correction(xyz, centers, weights, matrices, 1.3, kernel_type)) + torch.testing.assert_close(gradient, evaluate_micro_gradient(xyz, centers, weights, matrices, 1.3, kernel_type)) + assert no_gradient is None + assert torch.isfinite(gradient).all() + values[1].backward() + torch.testing.assert_close(xyz.grad[1], gradient[1]) + assert weights.grad is not None + + # ---------------------------------------------------------------- # evaluate_micro_correction # ---------------------------------------------------------------- diff --git a/tests/test_common/test_modules/test_evaluator/test_micro_anisotropic_macro_integration.py b/tests/test_common/test_modules/test_evaluator/test_micro_anisotropic_macro_integration.py index 982cdc13..4a771d59 100644 --- a/tests/test_common/test_modules/test_evaluator/test_micro_anisotropic_macro_integration.py +++ b/tests/test_common/test_modules/test_evaluator/test_micro_anisotropic_macro_integration.py @@ -16,10 +16,11 @@ surface_points_preprocess, orientations_preprocess, ) -from gempy_engine.modules.evaluator.micro_anisotropic_evaluator import ( +from .micro_reference import ( compute_anisotropy_matrices_from_gradients, solve_micro_weights, ) +from gempy_engine.modules.evaluator.micro_correction import MicroCorrection, apply_micro_correction PLOT = os.getenv("GEMPY_PLOT_MICRO", "0") == "1" @@ -34,10 +35,11 @@ def _build_grid_2d(x_range, y_range, nx, ny): return np.column_stack([xv.ravel(), yv.ravel()]) -def _eval_at_points(sp_internal, ori_internal, options, weights, xyz): +def _eval_at_points(sp_internal, ori_internal, options, weights, xyz, correction=None): eval_in = SolverInput(sp_internal, ori_internal, xyz_to_interpolate=xyz, fault_internal=None) options.evaluation_options.compute_scalar_gradient = True - return _evaluate_sys_eq(eval_in, weights, options) + fields = _evaluate_sys_eq(eval_in, weights, options) + return apply_micro_correction(fields, xyz, correction) @pytest.mark.skipif( @@ -124,14 +126,9 @@ def test_micro_correction_moves_contacts_closer_to_target(simple_model_2): options.evaluation_options.compute_scalar_gradient = False macro_fields = _eval_at_points(sp_internal, ori_internal, options, macro_weights, grid_xy) - micro = options.evaluation_options.micro_anisotropic - micro.enabled = True - micro.points = constraint_points - micro.weights = all_weights - micro.anisotropy_matrices = A - micro.kernel_range = micro_kernel_range + correction = MicroCorrection(constraint_points, A, all_weights, micro_kernel_range, "matern_5_2") - micro_fields = _eval_at_points(sp_internal, ori_internal, options, macro_weights, grid_xy) + micro_fields = _eval_at_points(sp_internal, ori_internal, options, macro_weights, grid_xy, correction) macro_field_2d = macro_fields.scalar_field.reshape(40, 40) micro_field_2d = micro_fields.scalar_field.reshape(40, 40) @@ -143,8 +140,7 @@ def test_micro_correction_moves_contacts_closer_to_target(simple_model_2): assert max_abs_diff > 1e-6, f"Micro correction should produce nonzero change, got max abs diff = {max_abs_diff}" # --- contact compliance --- - micro_exported = _eval_at_points(sp_internal, ori_internal, options, macro_weights, contacts) - options.evaluation_options.micro_anisotropic.enabled = False + micro_exported = _eval_at_points(sp_internal, ori_internal, options, macro_weights, contacts, correction) corrected_contacts = micro_exported.scalar_field rms_before = np.sqrt(np.mean(contact_residuals ** 2)) rms_after = np.sqrt(np.mean((target_values_at_contacts - corrected_contacts) ** 2)) @@ -155,9 +151,7 @@ def test_micro_correction_moves_contacts_closer_to_target(simple_model_2): # --- macro point preservation --- options.evaluation_options.compute_scalar_gradient = False - micro.enabled = True # re-enable for this eval - macro_after_exported = _eval_at_points(sp_internal, ori_internal, options, macro_weights, macro_sp_coords) - micro.enabled = False + macro_after_exported = _eval_at_points(sp_internal, ori_internal, options, macro_weights, macro_sp_coords, correction) macro_after_sp = macro_after_exported.scalar_field macro_drift = np.abs(macro_after_sp - macro_at_sp) max_macro_drift = np.max(macro_drift) @@ -376,10 +370,11 @@ def test_micro_correction_moves_3d_contacts_closer_to_target(simple_model): macro_weights = _solve_interpolation(solver_input, options.kernel_options) # --- target scalars: median macro scalar at original surface points --- - def _eval_3d(xyz): + def _eval_3d(xyz, correction=None): proxy = SolverInput(sp_internal, ori_internal, xyz_to_interpolate=xyz, fault_internal=None) options.evaluation_options.compute_scalar_gradient = True - return _evaluate_sys_eq(proxy, macro_weights, options) + fields = _evaluate_sys_eq(proxy, macro_weights, options) + return apply_micro_correction(fields, xyz, correction) exported_macro_sp = _eval_3d(macro_sp_coords) macro_at_sp = exported_macro_sp.scalar_field @@ -419,7 +414,7 @@ def _eval_3d(xyz): # --- build micro constraint system --- # When preserve_macro_points=True: contacts + macro SP as zero-residual constraints. # When False: contacts only. - micro = options.evaluation_options.micro_anisotropic + micro = options.micro_options micro.preserve_macro_points = False # this need to be false no question preserve = micro.preserve_macro_points @@ -458,14 +453,9 @@ def _eval_3d(xyz): options.evaluation_options.compute_scalar_gradient = False macro_fields = _eval_3d(grid_xyz) - micro = options.evaluation_options.micro_anisotropic - micro.enabled = True - micro.points = constraint_points - micro.weights = all_weights - micro.anisotropy_matrices = A - micro.kernel_range = micro_kernel_range + correction = MicroCorrection(constraint_points, A, all_weights, micro_kernel_range, "exponential") - micro_fields = _eval_3d(grid_xyz) + micro_fields = _eval_3d(grid_xyz, correction) macro_field = macro_fields.scalar_field micro_field = micro_fields.scalar_field @@ -477,9 +467,7 @@ def _eval_3d(xyz): assert max_abs_diff > 1e-8, f"3D micro correction should produce nonzero change, got max abs diff = {max_abs_diff}" # --- contact compliance --- - micro.enabled = True - micro_exported = _eval_3d(contacts) - micro.enabled = False + micro_exported = _eval_3d(contacts, correction) corrected_contacts = micro_exported.scalar_field rms_before = np.sqrt(np.mean(contact_residuals ** 2)) rms_after = np.sqrt(np.mean((target_values_at_contacts - corrected_contacts) ** 2)) @@ -491,9 +479,7 @@ def _eval_3d(xyz): # --- macro point preservation (only when macro SP are constraints) --- if preserve: options.evaluation_options.compute_scalar_gradient = False - micro.enabled = True - macro_after_exported = _eval_3d(macro_sp_coords) - micro.enabled = False + macro_after_exported = _eval_3d(macro_sp_coords, correction) macro_after_sp = macro_after_exported.scalar_field macro_drift = np.abs(macro_after_sp - macro_at_sp) max_macro_drift = np.max(macro_drift)