Skip to content

Commit cb40625

Browse files
domfournierdomfournier
committed
Move utilities
1 parent d1e5495 commit cb40625

4 files changed

Lines changed: 135 additions & 161 deletions

File tree

‎simpeg_drivers/components/data.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -398,8 +398,7 @@ def update_params(self, data_dict, uncert_dict):
398398
if getattr(self.params, "line_selection", None) is not None:
399399
if self.params.line_selection.line_object is None:
400400
parts = get_parts_from_electrodes(self.entity)
401-
_, u_part = np.unique(parts, return_inverse=True)
402-
line_ids = self.entity.add_data({"Line IDs": {"values": u_part + 1}})
401+
line_ids = self.entity.add_data({"Line IDs": {"values": parts + 1}})
403402
else:
404403
line_ids = self.params.line_selection.line_object.copy(
405404
parent=self.entity,

‎simpeg_drivers/electricals/base_2d.py‎

Lines changed: 1 addition & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -11,19 +11,11 @@
1111

1212
from __future__ import annotations
1313

14-
import numpy as np
15-
from geoh5py.data import IntegerData
16-
from geoh5py.objects import DrapeModel
17-
from geoh5py.shared.merging.drape_model import DrapeModelMerger
1814
from geoh5py.ui_json.ui_json import fetch_active_workspace
19-
from geoh5py.workspace import Workspace
2015

2116
from simpeg_drivers.components.meshes import InversionMesh
2217
from simpeg_drivers.driver import InversionDriver
23-
from simpeg_drivers.options import (
24-
DrapeModelOptions,
25-
)
26-
from simpeg_drivers.utils.utils import get_drape_model
18+
from simpeg_drivers.utils.surveys import create_mesh_by_line_id
2719

2820

2921
class Base2DDriver(InversionDriver):
@@ -53,102 +45,3 @@ def inversion_mesh(self) -> InversionMesh:
5345
)
5446

5547
return self._inversion_mesh
56-
57-
58-
def create_mesh_by_line_id(
59-
workspace: Workspace,
60-
line_ids: IntegerData,
61-
drape_options: DrapeModelOptions,
62-
**object_kwargs,
63-
) -> DrapeModel:
64-
"""
65-
Create a drape mesh for the dc resistivity survey lines.
66-
67-
:param workspace: Workspace to create the drape mesh in.
68-
:param line_ids: IntegerData object containing the line IDs for each vertex.
69-
:param drape_options: DrapeModelOptions containing the parameters for the drape mesh
70-
:param object_kwargs: Additional keyword arguments to pass to the DrapeModelMerger.create_object method.
71-
72-
:return: A DrapeModel object containing the merged drape mesh for all survey lines.
73-
"""
74-
drape_models = []
75-
temp_work = Workspace()
76-
77-
relief = get_max_line_relief(line_ids, drape_options.v_cell_size)
78-
79-
for line_id in np.unique(line_ids.values):
80-
poles = get_poles_by_line_id(line_ids, line_id)
81-
poles = np.unique(poles, axis=0)
82-
poles = normalize_vertically(poles, relief)
83-
84-
drape_model = get_drape_model(
85-
temp_work,
86-
poles,
87-
[
88-
drape_options.u_cell_size,
89-
drape_options.v_cell_size,
90-
],
91-
drape_options.depth_core,
92-
[drape_options.horizontal_padding] * 2
93-
+ [drape_options.vertical_padding, 1],
94-
drape_options.expansion_factor,
95-
)
96-
drape_models.append(drape_model)
97-
98-
entity = DrapeModelMerger.create_object(workspace, drape_models, **object_kwargs)
99-
100-
return entity
101-
102-
103-
def get_max_line_relief(line_ids, z_cell_size) -> float:
104-
"""
105-
Get the maximum relief across all survey lines, rounded to the nearest cell thickness.
106-
107-
:param line_ids: IntegerData object containing the line IDs for each vertex.
108-
:param z_cell_size: Cell size in the vertical direction for the drape mesh.
109-
"""
110-
max_relief = 0
111-
for line_id in np.unique(line_ids.values):
112-
poles = get_poles_by_line_id(line_ids, line_id)
113-
max_relief = np.maximum(poles[:, 2].max() - poles[:, 2].min(), max_relief)
114-
115-
return (max_relief // z_cell_size + 2) * z_cell_size
116-
117-
118-
def normalize_vertically(poles: np.ndarray, relief: float) -> np.ndarray:
119-
"""
120-
Given a set of pole locations, normalize the vertical component to the maximum relief across all lines.
121-
122-
This ensures that the drape mesh has uniform vertical discretization across all survey lines.
123-
124-
:param poles: Array of pole locations to normalize.
125-
:param relief: Maximum relief across all survey lines, rounded to the nearest cell thickness.
126-
127-
:return: Array of pole locations with normalized vertical component.
128-
"""
129-
min_poles_z = poles[:, 2].min()
130-
poles[:, 2] -= min_poles_z
131-
poles[:, 2] *= relief / np.maximum(poles[:, 2].max(), 1e-3)
132-
133-
# Shift back vertically
134-
poles[:, 2] += min_poles_z
135-
136-
return poles
137-
138-
139-
def get_poles_by_line_id(line_ids: IntegerData, uid: int) -> np.ndarray:
140-
"""Get the vertices associated with a given line ID."""
141-
mn_mask = line_ids.values == uid
142-
143-
unique_tx = np.unique(line_ids.parent.ab_cell_id.values[mn_mask])
144-
145-
ab_mask = np.isin(line_ids.parent.complement.ab_cell_id.values, unique_tx)
146-
147-
return np.vstack(
148-
[
149-
line_ids.parent.vertices[line_ids.parent.cells[mn_mask].flatten()],
150-
line_ids.parent.current_electrodes.vertices[
151-
line_ids.parent.current_electrodes.cells[ab_mask].flatten()
152-
],
153-
]
154-
)

‎simpeg_drivers/utils/surveys.py‎

Lines changed: 110 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,19 @@
1313

1414
import numpy as np
1515
from discretize import TreeMesh
16-
from geoapps_utils.utils.numerical import traveling_salesman
1716
from geoh5py import Workspace
18-
from geoh5py.objects import PotentialElectrode
17+
from geoh5py.data import IntegerData
18+
from geoh5py.objects import DrapeModel, PotentialElectrode
19+
from geoh5py.shared.merging.drape_model import DrapeModelMerger
1920
from scipy.sparse import csgraph, csr_matrix
2021
from scipy.spatial import cKDTree
2122
from simpeg.survey import BaseSurvey
2223

24+
from simpeg_drivers.options import (
25+
DrapeModelOptions,
26+
)
27+
from simpeg_drivers.utils.utils import get_drape_model
28+
2329

2430
def station_spacing(
2531
locations: np.ndarray,
@@ -63,47 +69,6 @@ def counter_clockwise_sort(segments: np.ndarray, vertices: np.ndarray) -> np.nda
6369
return segments
6470

6571

66-
def compute_alongline_distance(points: np.ndarray, ordered: bool = True):
67-
"""
68-
Convert from cartesian (x, y, values) points to (distance, values) locations.
69-
70-
:param: points: Cartesian coordinates of points lying either roughly within a
71-
plane or a line.
72-
"""
73-
if not ordered:
74-
order = traveling_salesman(points)
75-
points = points[order, :]
76-
77-
distances = np.cumsum(
78-
np.r_[0, np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1)]
79-
)
80-
if points.shape[1] > 2:
81-
distances = np.c_[distances, points[:, 2:]]
82-
83-
return distances
84-
85-
86-
def copy_potentials_from_mask(
87-
workspace: Workspace, survey: PotentialElectrode, cell_mask: np.ndarray
88-
):
89-
"""
90-
Returns a survey containing data from a single line.
91-
92-
:param workspace: geoh5py workspace containing a valid DCIP survey.
93-
:param survey: PotentialElectrode object.
94-
:param cell_mask: Boolean array of M-N pairs to include in the new survey.
95-
"""
96-
97-
if not np.any(cell_mask):
98-
raise ValueError("No cells found in the mask.")
99-
100-
active_poles = np.zeros(survey.n_vertices, dtype=bool)
101-
active_poles[survey.cells[cell_mask, :].ravel()] = True
102-
potentials = survey.copy(parent=workspace, mask=active_poles, cell_mask=cell_mask)
103-
104-
return potentials
105-
106-
10772
def get_intersecting_cells(locations: np.ndarray, mesh: TreeMesh) -> np.ndarray:
10873
"""
10974
Find cells that intersect with a set of segments.
@@ -166,7 +131,9 @@ def get_parts_from_electrodes(survey: PotentialElectrode) -> np.ndarray:
166131
)
167132

168133
connections = csgraph.connected_components(edge_array)[1]
169-
return connections[survey.cells[:, 0]]
134+
parts = connections[survey.cells[:, 0]]
135+
_, u_part = np.unique(parts, return_inverse=True)
136+
return u_part
170137

171138

172139
def compute_em_projections(locations, simulation):
@@ -206,3 +173,102 @@ def compute_dc_projections(locations, cells, simulation):
206173
proj_mn -= projection[cells[indices, 1], :]
207174

208175
receiver.spatialP = proj_mn # pylint: disable=protected-access
176+
177+
178+
def create_mesh_by_line_id(
179+
workspace: Workspace,
180+
line_ids: IntegerData,
181+
drape_options: DrapeModelOptions,
182+
**object_kwargs,
183+
) -> DrapeModel:
184+
"""
185+
Create a drape mesh for the dc resistivity survey lines.
186+
187+
:param workspace: Workspace to create the drape mesh in.
188+
:param line_ids: IntegerData object containing the line IDs for each vertex.
189+
:param drape_options: DrapeModelOptions containing the parameters for the drape mesh
190+
:param object_kwargs: Additional keyword arguments to pass to the DrapeModelMerger.create_object method.
191+
192+
:return: A DrapeModel object containing the merged drape mesh for all survey lines.
193+
"""
194+
drape_models = []
195+
temp_work = Workspace()
196+
197+
relief = get_max_line_relief(line_ids, drape_options.v_cell_size)
198+
199+
for line_id in np.unique(line_ids.values):
200+
poles = get_poles_by_line_id(line_ids, line_id)
201+
poles = np.unique(poles, axis=0)
202+
poles = normalize_vertically(poles, relief)
203+
204+
drape_model = get_drape_model(
205+
temp_work,
206+
poles,
207+
[
208+
drape_options.u_cell_size,
209+
drape_options.v_cell_size,
210+
],
211+
drape_options.depth_core,
212+
[drape_options.horizontal_padding] * 2
213+
+ [drape_options.vertical_padding, 1],
214+
drape_options.expansion_factor,
215+
)
216+
drape_models.append(drape_model)
217+
218+
entity = DrapeModelMerger.create_object(workspace, drape_models, **object_kwargs)
219+
220+
return entity
221+
222+
223+
def get_max_line_relief(line_ids: IntegerData, z_cell_size: float) -> float:
224+
"""
225+
Get the maximum relief across all survey lines, rounded to the nearest cell thickness.
226+
227+
:param line_ids: IntegerData object containing the line IDs for each vertex.
228+
:param z_cell_size: Cell size in the vertical direction for the drape mesh.
229+
"""
230+
max_relief = 0
231+
for line_id in np.unique(line_ids.values):
232+
poles = get_poles_by_line_id(line_ids, line_id)
233+
max_relief = np.maximum(poles[:, 2].max() - poles[:, 2].min(), max_relief)
234+
235+
return (max_relief // z_cell_size + 2) * z_cell_size
236+
237+
238+
def get_poles_by_line_id(line_ids: IntegerData, uid: int) -> np.ndarray:
239+
"""Get the vertices associated with a given line ID."""
240+
mn_mask = line_ids.values == uid
241+
242+
unique_tx = np.unique(line_ids.parent.ab_cell_id.values[mn_mask])
243+
244+
ab_mask = np.isin(line_ids.parent.complement.ab_cell_id.values, unique_tx)
245+
246+
return np.vstack(
247+
[
248+
line_ids.parent.vertices[line_ids.parent.cells[mn_mask].flatten()],
249+
line_ids.parent.current_electrodes.vertices[
250+
line_ids.parent.current_electrodes.cells[ab_mask].flatten()
251+
],
252+
]
253+
)
254+
255+
256+
def normalize_vertically(poles: np.ndarray, relief: float) -> np.ndarray:
257+
"""
258+
Given a set of pole locations, normalize the vertical component to the maximum relief across all lines.
259+
260+
This ensures that the drape mesh has uniform vertical discretization across all survey lines.
261+
262+
:param poles: Array of pole locations to normalize.
263+
:param relief: Maximum relief across all survey lines, rounded to the nearest cell thickness.
264+
265+
:return: Array of pole locations with normalized vertical component.
266+
"""
267+
min_poles_z = poles[:, 2].min()
268+
poles[:, 2] -= min_poles_z
269+
poles[:, 2] *= relief / np.maximum(poles[:, 2].max(), 1e-3)
270+
271+
# Shift back vertically
272+
poles[:, 2] += min_poles_z
273+
274+
return poles

‎simpeg_drivers/utils/utils.py‎

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
from geoapps_utils.utils.numerical import running_mean, traveling_salesman
2323
from geoh5py import Workspace
2424
from geoh5py.data import NumericData
25-
from geoh5py.groups import Group, SimPEGGroup
25+
from geoh5py.groups import SimPEGGroup
2626
from geoh5py.objects import DrapeModel, Octree
2727
from geoh5py.objects.surveys.direct_current import PotentialElectrode
2828
from geoh5py.objects.surveys.electromagnetics.airborne_app_con import (
@@ -32,14 +32,10 @@
3232
from geoh5py.shared import INTEGER_NDV
3333
from geoh5py.ui_json import InputFile
3434
from grid_apps.utils import octree_2_treemesh
35-
from scipy.interpolate import LinearNDInterpolator, NearestNDInterpolator, interp1d
36-
from scipy.sparse import csr_matrix, diags
37-
from scipy.spatial import ConvexHull, Delaunay, cKDTree
35+
from scipy.interpolate import interp1d
36+
from scipy.spatial import ConvexHull, cKDTree
3837

3938
from simpeg_drivers import DRIVER_MAP
40-
from simpeg_drivers.utils.surveys import (
41-
compute_alongline_distance,
42-
)
4339

4440

4541
if TYPE_CHECKING:
@@ -629,3 +625,23 @@ def simpeg_group_to_driver(group: SimPEGGroup, workspace: Workspace) -> Inversio
629625
params = inversion_driver._params_class.build(ifile) # pylint: disable=protected-access
630626

631627
return inversion_driver(params)
628+
629+
630+
def compute_alongline_distance(points: np.ndarray, ordered: bool = True):
631+
"""
632+
Convert from cartesian (x, y, values) points to (distance, values) locations.
633+
634+
:param: points: Cartesian coordinates of points lying either roughly within a
635+
plane or a line.
636+
"""
637+
if not ordered:
638+
order = traveling_salesman(points)
639+
points = points[order, :]
640+
641+
distances = np.cumsum(
642+
np.r_[0, np.linalg.norm(np.diff(points[:, :2], axis=0), axis=1)]
643+
)
644+
if points.shape[1] > 2:
645+
distances = np.c_[distances, points[:, 2:]]
646+
647+
return distances

0 commit comments

Comments
 (0)