interpolate.py (13170B)
1 import math 2 from collections import namedtuple 3 4 import numpy as np 5 6 from .core import resistivity_color 7 8 VOXEL_SAFETY_CAP = 5_000_000 9 10 Sample = namedtuple('Sample', ['x', 'y', 'z', 'resistivity', 'doi_elev']) 11 12 13 def axis_count(extent, cell_size): 14 if cell_size <= 0: 15 raise ValueError('cell size must be positive') 16 if extent <= 0: 17 raise ValueError('extent must be positive') 18 return max(1, math.ceil(extent / cell_size)) 19 20 21 def axis_cell_size(extent, count): 22 if count < 1: 23 raise ValueError('cell count must be >= 1') 24 if extent <= 0: 25 raise ValueError('extent must be positive') 26 return extent / count 27 28 29 class GridSpec: 30 def __init__(self, xmin, ymin, zmin, dx, dy, dz, nx, ny, nz): 31 self.xmin = xmin 32 self.ymin = ymin 33 self.zmin = zmin 34 self.dx = dx 35 self.dy = dy 36 self.dz = dz 37 self.nx = nx 38 self.ny = ny 39 self.nz = nz 40 41 @classmethod 42 def from_cell_size(cls, bounds, dx, dy, dz): 43 xmin, xmax, ymin, ymax, zmin, zmax = bounds 44 nx = axis_count(xmax - xmin, dx) 45 ny = axis_count(ymax - ymin, dy) 46 nz = axis_count(zmax - zmin, dz) 47 return cls(xmin, ymin, zmin, dx, dy, dz, nx, ny, nz) 48 49 @classmethod 50 def from_cell_count(cls, bounds, nx, ny, nz): 51 xmin, xmax, ymin, ymax, zmin, zmax = bounds 52 dx = axis_cell_size(xmax - xmin, nx) 53 dy = axis_cell_size(ymax - ymin, ny) 54 dz = axis_cell_size(zmax - zmin, nz) 55 return cls(xmin, ymin, zmin, dx, dy, dz, nx, ny, nz) 56 57 def cell_count(self): 58 return self.nx * self.ny * self.nz 59 60 def centers(self): 61 for iz in range(self.nz): 62 z = self.zmin + (iz + 0.5) * self.dz 63 for iy in range(self.ny): 64 y = self.ymin + (iy + 0.5) * self.dy 65 for ix in range(self.nx): 66 x = self.xmin + (ix + 0.5) * self.dx 67 yield ix, iy, iz, x, y, z 68 69 70 def doi_by_xy_from_points(doi_points): 71 result = {} 72 for pt in doi_points: 73 x = pt.get('X') if isinstance(pt, dict) else pt[0] 74 y = pt.get('Y') if isinstance(pt, dict) else pt[1] 75 zdoi = pt.get('ZDOI') if isinstance(pt, dict) else pt[2] 76 if x is not None and y is not None and zdoi is not None: 77 result[(float(x), float(y))] = float(zdoi) 78 return result 79 80 81 def samples_from_layers(layer_rows, doi_by_xy=None): 82 samples = [] 83 for row in layer_rows: 84 try: 85 x = float(row['X']) 86 y = float(row['Y']) 87 z = float(row['ZMid']) 88 except (KeyError, ValueError, TypeError): 89 continue 90 if math.isnan(x) or math.isnan(y) or math.isnan(z): 91 continue 92 res = row.get('Resistivity') 93 if res is None: 94 continue 95 try: 96 res = float(res) 97 except (ValueError, TypeError): 98 continue 99 if math.isnan(res): 100 continue 101 doi_elev = None 102 if doi_by_xy is not None: 103 # Joined by exact (x, y); both layers come from the same loader 104 # parse, so their coordinates are bit-identical. 105 doi_elev = doi_by_xy.get((x, y)) 106 samples.append(Sample(x, y, z, res, doi_elev)) 107 return samples 108 109 110 class _SampleArrays: 111 """Columnar view of the sample set for vectorized interpolation. 112 113 Missing DOI elevations are stored as NaN so the gate test reduces to a 114 single masked comparison.""" 115 116 __slots__ = ('x', 'y', 'z', 'res', 'doi') 117 118 def __init__(self, samples): 119 n = len(samples) 120 self.x = np.fromiter((s.x for s in samples), dtype=float, count=n) 121 self.y = np.fromiter((s.y for s in samples), dtype=float, count=n) 122 self.z = np.fromiter((s.z for s in samples), dtype=float, count=n) 123 self.res = np.fromiter((s.resistivity for s in samples), dtype=float, 124 count=n) 125 self.doi = np.fromiter( 126 (np.nan if s.doi_elev is None else s.doi_elev for s in samples), 127 dtype=float, count=n, 128 ) 129 130 def __len__(self): 131 return self.x.shape[0] 132 133 134 def _default_radius(samples): 135 """Three times the median nearest-neighbor spacing of the distinct sounding 136 locations. Operating on unique (x, y) positions makes the estimate 137 independent of feature ordering and of how many depth layers each sounding 138 has.""" 139 locations = sorted({(s.x, s.y) for s in samples}) 140 if len(locations) < 2: 141 return None 142 # The nearest-neighbor scan is O(n**2); cap the working set with an 143 # evenly-spaced subsample so dense surveys stay fast (the subsample slightly 144 # over-estimates spacing, which the ample 3x multiplier absorbs). 145 cap = 500 146 if len(locations) > cap: 147 step = len(locations) / cap 148 locations = [locations[int(i * step)] for i in range(cap)] 149 pts = np.asarray(locations, dtype=float) 150 diff = pts[:, None, :] - pts[None, :, :] 151 dist = np.hypot(diff[..., 0], diff[..., 1]) 152 np.fill_diagonal(dist, np.inf) 153 nearest = dist.min(axis=1) 154 nearest = nearest[nearest > 0] 155 if nearest.size == 0: 156 return None 157 return float(np.median(nearest) * 3.0) 158 159 160 def _default_radius_z(samples): 161 """Vertical search window derived from the typical layer thickness: 1.5x the 162 median gap between consecutive sample depths within a single sounding 163 (roughly a couple of layers). Operating per sounding keeps the estimate 164 independent of how the soundings are laid out horizontally. Returns ``None`` 165 when no sounding has two or more distinct depths.""" 166 by_xy = {} 167 for s in samples: 168 by_xy.setdefault((s.x, s.y), []).append(s.z) 169 gaps = [] 170 for zs in by_xy.values(): 171 if len(zs) < 2: 172 continue 173 zs = sorted(zs) 174 for a, b in zip(zs, zs[1:]): 175 gap = b - a 176 if gap > 0: 177 gaps.append(gap) 178 if not gaps: 179 return None 180 return float(np.median(gaps) * 1.5) 181 182 183 def _bucket_index(arrays, radius): 184 """Map integer (bx, by) cells of side ``radius`` to the sample indices that 185 fall inside them, so a query only weighs samples in its 3x3 neighborhood.""" 186 if radius is None or radius <= 0: 187 return None, None 188 bx = np.floor(arrays.x / radius).astype(np.int64) 189 by = np.floor(arrays.y / radius).astype(np.int64) 190 buckets = {} 191 for i in range(len(arrays)): 192 key = (int(bx[i]), int(by[i])) 193 buckets.setdefault(key, []).append(i) 194 return {k: np.asarray(v) for k, v in buckets.items()}, radius 195 196 197 def _candidate_indices(buckets, bucket_size, x, y): 198 if buckets is None: 199 return None 200 bx = int(math.floor(x / bucket_size)) 201 by = int(math.floor(y / bucket_size)) 202 parts = [] 203 for dbx in (-1, 0, 1): 204 for dby in (-1, 0, 1): 205 arr = buckets.get((bx + dbx, by + dby)) 206 if arr is not None: 207 parts.append(arr) 208 if not parts: 209 return np.empty(0, dtype=np.int64) 210 return np.concatenate(parts) 211 212 213 def _interpolate_at(arrays, idx, qx, qy, qz, *, power, radius, radius_z, 214 max_neighbors, min_neighbors, log10, gate_doi): 215 """Inverse-distance interpolation at a single query point. 216 217 ``radius`` (horizontal) and ``radius_z`` (vertical) define a search 218 ellipsoid that governs both which samples are used and how they are 219 weighted. The ellipsoid is a sphere in vertically-scaled coordinates with 220 exaggeration ``k = radius / radius_z``, so the scaled squared distance 221 ``d2 = horizontal2 + (dz * k)**2`` drives the cull, neighbor selection, and 222 weighting consistently. When either radius is ``None`` weighting is isotropic 223 (``k = 1``) and only the axes whose radius is set are culled. 224 225 Returns ``(value, num_valid)``. ``value`` is ``None`` when fewer than 226 ``min_neighbors`` usable samples remain; ``num_valid`` is always the count of 227 samples that passed the DOI gate and ellipsoid cull. Shared by :func:`idw` 228 and :func:`interpolate_voxels` so both honor identical rules.""" 229 if idx is None: 230 sx, sy, sz, sres, sdoi = (arrays.x, arrays.y, arrays.z, arrays.res, 231 arrays.doi) 232 else: 233 sx, sy, sz, sres, sdoi = (arrays.x[idx], arrays.y[idx], arrays.z[idx], 234 arrays.res[idx], arrays.doi[idx]) 235 236 dx = sx - qx 237 dy = sy - qy 238 dz = sz - qz 239 horizontal2 = dx * dx + dy * dy 240 241 if radius is not None and radius_z is not None: 242 k = radius / radius_z 243 else: 244 k = 1.0 245 vdz = dz * k 246 d2 = horizontal2 + vdz * vdz 247 248 keep = np.ones(horizontal2.shape, dtype=bool) 249 if gate_doi: 250 # NaN DOI (no DOI for that sounding) never gates; only drop samples 251 # whose DOI elevation sits above the query depth. 252 keep &= ~((~np.isnan(sdoi)) & (qz < sdoi)) 253 if radius is not None and radius_z is not None: 254 # Ellipsoid: d2 <= radius**2 is horizontal2/radius**2 + dz**2/radius_z**2 <= 1. 255 keep &= d2 <= radius * radius 256 else: 257 if radius is not None: 258 keep &= horizontal2 <= radius * radius 259 if radius_z is not None: 260 keep &= dz * dz <= radius_z * radius_z 261 262 d2 = d2[keep] 263 res = sres[keep] 264 265 if max_neighbors is not None and d2.shape[0] > max_neighbors: 266 nearest = np.argpartition(d2, max_neighbors)[:max_neighbors] 267 d2 = d2[nearest] 268 res = res[nearest] 269 270 num_valid = int(d2.shape[0]) 271 if num_valid < min_neighbors: 272 return None, num_valid 273 274 # Coincident samples: average all exact hits rather than picking the first. 275 exact = res[d2 == 0.0] 276 if exact.size: 277 usable = exact[exact > 0] if log10 else exact 278 if usable.size: 279 if log10: 280 return float(10.0 ** np.log10(usable).mean()), num_valid 281 return float(usable.mean()), num_valid 282 283 if log10: 284 # In log space a non-positive resistivity has no defined log; drop it 285 # entirely rather than letting it contribute weight at log10(res)==0. 286 positive = res > 0 287 d2 = d2[positive] 288 res = res[positive] 289 290 if d2.size == 0: 291 return None, num_valid 292 293 weights = 1.0 / (d2 ** (power / 2.0)) 294 total_w = float(weights.sum()) 295 if total_w == 0.0: 296 return None, num_valid 297 298 values = np.log10(res) if log10 else res 299 avg = float((weights * values).sum() / total_w) 300 return (10.0 ** avg if log10 else avg), num_valid 301 302 303 def idw(samples, queries, power=2.0, radius=None, radius_z=None, 304 max_neighbors=None, min_neighbors=1, log10=True, gate_doi=True): 305 if radius is None: 306 radius = _default_radius(samples) 307 308 arrays = _SampleArrays(samples) 309 buckets, bucket_size = _bucket_index(arrays, radius) 310 311 results = [] 312 for qx, qy, qz in queries: 313 idx = _candidate_indices(buckets, bucket_size, qx, qy) 314 value, _ = _interpolate_at( 315 arrays, idx, qx, qy, qz, 316 power=power, radius=radius, radius_z=radius_z, 317 max_neighbors=max_neighbors, min_neighbors=min_neighbors, 318 log10=log10, gate_doi=gate_doi, 319 ) 320 results.append(value) 321 return results 322 323 324 def interpolate_voxels(layer_rows, grid, *, doi_points=None, 325 clip_below_doi=True, power=2.0, radius=None, 326 radius_z=None, max_neighbors=None, min_neighbors=1, 327 log10=True, progress=None): 328 if grid.cell_count() > VOXEL_SAFETY_CAP: 329 raise ValueError( 330 f'grid has {grid.cell_count()} cells, exceeding the safety cap of ' 331 f'{VOXEL_SAFETY_CAP}; reduce the cell count or increase cell size' 332 ) 333 334 doi_by_xy = None 335 if doi_points is not None: 336 doi_by_xy = doi_by_xy_from_points(doi_points) 337 338 samples = samples_from_layers(layer_rows, doi_by_xy) 339 340 if radius is None: 341 radius = _default_radius(samples) 342 if radius_z is None: 343 radius_z = _default_radius_z(samples) 344 345 arrays = _SampleArrays(samples) 346 buckets, bucket_size = _bucket_index(arrays, radius) 347 348 voxels = [] 349 total = grid.cell_count() 350 done = 0 351 352 for ix, iy, iz, cx, cy, cz in grid.centers(): 353 idx = _candidate_indices(buckets, bucket_size, cx, cy) 354 res_val, num_valid = _interpolate_at( 355 arrays, idx, cx, cy, cz, 356 power=power, radius=radius, radius_z=radius_z, 357 max_neighbors=max_neighbors, min_neighbors=min_neighbors, 358 log10=log10, gate_doi=clip_below_doi, 359 ) 360 done += 1 361 362 if res_val is not None: 363 voxels.append({ 364 'X': cx, 365 'Y': cy, 366 'Z': cz, 367 'Resistivity': res_val, 368 'Color': resistivity_color(res_val), 369 'CellX': ix, 370 'CellY': iy, 371 'CellZ': iz, 372 'NumSamples': num_valid, 373 'Geometry': f'POINT Z ({cx} {cy} {cz})', 374 }) 375 376 # A truthy return from progress requests cancellation. 377 if progress is not None and progress(done, total): 378 break 379 380 return voxels