test_interpolate.py (17270B)
1 import math 2 import unittest 3 4 from tem_loader.interpolate import ( 5 GridSpec, 6 Sample, 7 _default_radius, 8 _default_radius_z, 9 axis_cell_size, 10 axis_count, 11 doi_by_xy_from_points, 12 idw, 13 interpolate_voxels, 14 samples_from_layers, 15 ) 16 17 18 class AxisMathTests(unittest.TestCase): 19 def test_axis_count_ceil(self): 20 self.assertEqual(axis_count(10.0, 3.0), 4) 21 self.assertEqual(axis_count(9.0, 3.0), 3) 22 self.assertEqual(axis_count(0.1, 3.0), 1) 23 24 def test_axis_count_exact(self): 25 self.assertEqual(axis_count(100.0, 10.0), 10) 26 27 def test_axis_count_rejects_non_positive_cell_size(self): 28 with self.assertRaises(ValueError): 29 axis_count(10.0, 0.0) 30 with self.assertRaises(ValueError): 31 axis_count(10.0, -1.0) 32 33 def test_axis_count_rejects_non_positive_extent(self): 34 with self.assertRaises(ValueError): 35 axis_count(0.0, 1.0) 36 with self.assertRaises(ValueError): 37 axis_count(-5.0, 1.0) 38 39 def test_axis_cell_size_round_trip(self): 40 extent = 100.0 41 cell_size = 7.0 42 n = axis_count(extent, cell_size) 43 cs = axis_cell_size(extent, n) 44 self.assertLessEqual(cs, cell_size) 45 self.assertGreaterEqual(n * cs, extent) 46 47 def test_axis_cell_size_rejects_zero_count(self): 48 with self.assertRaises(ValueError): 49 axis_cell_size(10.0, 0) 50 51 def test_axis_cell_size_rejects_non_positive_extent(self): 52 with self.assertRaises(ValueError): 53 axis_cell_size(0.0, 5) 54 55 56 class GridSpecTests(unittest.TestCase): 57 def test_from_cell_size_cell_count(self): 58 bounds = (0.0, 100.0, 0.0, 50.0, -10.0, 0.0) 59 g = GridSpec.from_cell_size(bounds, 10.0, 10.0, 2.0) 60 self.assertEqual(g.nx, 10) 61 self.assertEqual(g.ny, 5) 62 self.assertEqual(g.nz, 5) 63 self.assertEqual(g.cell_count(), 250) 64 65 def test_from_cell_count_consistency(self): 66 bounds = (0.0, 100.0, 0.0, 50.0, -10.0, 0.0) 67 g = GridSpec.from_cell_count(bounds, 5, 5, 2) 68 self.assertEqual(g.nx, 5) 69 self.assertEqual(g.ny, 5) 70 self.assertEqual(g.nz, 2) 71 self.assertAlmostEqual(g.dx, 20.0) 72 self.assertAlmostEqual(g.dy, 10.0) 73 self.assertAlmostEqual(g.dz, 5.0) 74 75 def test_centers_half_cell_offset(self): 76 g = GridSpec(0.0, 0.0, 0.0, 10.0, 10.0, 10.0, 2, 2, 1) 77 centers = list(g.centers()) 78 self.assertEqual(len(centers), 4) 79 xs = sorted({c[3] for c in centers}) 80 ys = sorted({c[4] for c in centers}) 81 zs = sorted({c[5] for c in centers}) 82 self.assertAlmostEqual(xs[0], 5.0) 83 self.assertAlmostEqual(xs[1], 15.0) 84 self.assertAlmostEqual(ys[0], 5.0) 85 self.assertAlmostEqual(ys[1], 15.0) 86 self.assertAlmostEqual(zs[0], 5.0) 87 88 def test_covered_extent_at_least_requested(self): 89 bounds = (0.0, 7.0, 0.0, 5.0, 0.0, 3.0) 90 g = GridSpec.from_cell_size(bounds, 3.0, 3.0, 3.0) 91 self.assertGreaterEqual(g.nx * g.dx, 7.0) 92 self.assertGreaterEqual(g.ny * g.dy, 5.0) 93 self.assertGreaterEqual(g.nz * g.dz, 3.0) 94 95 96 class IDWTests(unittest.TestCase): 97 def _s(self, x, y, z, res, doi=None): 98 return Sample(x, y, z, res, doi) 99 100 def test_single_sample_returns_its_value(self): 101 samples = [self._s(0.0, 0.0, 0.0, 100.0)] 102 result = idw(samples, [(1.0, 0.0, 0.0)], radius=None) 103 self.assertAlmostEqual(result[0], 100.0, places=3) 104 105 def test_equidistant_power2_is_mean_in_log_space(self): 106 samples = [ 107 self._s(-1.0, 0.0, 0.0, 10.0), 108 self._s(1.0, 0.0, 0.0, 1000.0), 109 ] 110 result = idw(samples, [(0.0, 0.0, 0.0)], power=2.0, log10=True, radius=None) 111 expected = 10.0 ** ((math.log10(10.0) + math.log10(1000.0)) / 2.0) 112 self.assertAlmostEqual(result[0], expected, places=5) 113 114 def test_exact_hit_returns_exact_value(self): 115 samples = [ 116 self._s(0.0, 0.0, 0.0, 50.0), 117 self._s(1.0, 0.0, 0.0, 200.0), 118 ] 119 result = idw(samples, [(0.0, 0.0, 0.0)], radius=None) 120 self.assertAlmostEqual(result[0], 50.0) 121 122 def test_radius_excludes_far_points(self): 123 samples = [ 124 self._s(0.0, 0.0, 0.0, 50.0), 125 self._s(100.0, 0.0, 0.0, 9999.0), 126 ] 127 result = idw(samples, [(1.0, 0.0, 0.0)], radius=10.0, min_neighbors=1) 128 self.assertIsNotNone(result[0]) 129 self.assertAlmostEqual(result[0], 50.0, places=1) 130 131 def test_min_neighbors_unmet_returns_none(self): 132 samples = [self._s(0.0, 0.0, 0.0, 50.0)] 133 result = idw(samples, [(100.0, 0.0, 0.0)], radius=5.0, min_neighbors=2) 134 self.assertIsNone(result[0]) 135 136 def test_radius_z_excludes_distant_depths(self): 137 # Sample A sits 5 units away horizontally; sample B sits 50 units away 138 # vertically. A horizontal radius of 100 keeps both, but a tight vertical 139 # radius of 10 culls B, leaving A's value only. 140 samples = [ 141 self._s(5.0, 0.0, 0.0, 10.0), 142 self._s(0.0, 0.0, 50.0, 100.0), 143 ] 144 with_both = idw(samples, [(0.0, 0.0, 0.0)], radius=100.0, radius_z=100.0, 145 log10=False) 146 tight_z = idw(samples, [(0.0, 0.0, 0.0)], radius=100.0, radius_z=10.0, 147 log10=False) 148 self.assertIsNotNone(with_both[0]) 149 self.assertAlmostEqual(tight_z[0], 10.0, places=6) 150 self.assertNotAlmostEqual(with_both[0], tight_z[0], places=1) 151 152 def test_log10_two_samples_geometric_mean(self): 153 samples = [ 154 self._s(-1.0, 0.0, 0.0, 10.0), 155 self._s(1.0, 0.0, 0.0, 1000.0), 156 ] 157 result = idw(samples, [(0.0, 0.0, 0.0)], log10=True, radius=None) 158 expected = math.sqrt(10.0 * 1000.0) 159 self.assertAlmostEqual(result[0], expected, places=4) 160 161 def test_doi_gate_drops_sounding_below_doi(self): 162 samples = [ 163 self._s(0.0, 0.0, -50.0, 100.0, doi=-30.0), 164 self._s(1.0, 0.0, -10.0, 200.0, doi=None), 165 ] 166 result = idw(samples, [(0.5, 0.0, -40.0)], gate_doi=True, radius=None, 167 log10=False) 168 result_no_gate = idw(samples, [(0.5, 0.0, -40.0)], gate_doi=False, 169 radius=None, log10=False) 170 self.assertIsNotNone(result[0]) 171 self.assertIsNotNone(result_no_gate[0]) 172 self.assertNotAlmostEqual(result[0], result_no_gate[0], places=1) 173 174 def test_doi_gate_none_doi_never_gated(self): 175 samples = [self._s(0.0, 0.0, -100.0, 50.0, doi=None)] 176 result = idw(samples, [(0.0, 0.0, -200.0)], gate_doi=True, radius=None) 177 self.assertIsNotNone(result[0]) 178 179 def test_log10_skips_non_positive_resistivity(self): 180 # A zero/negative resistivity has no log; it must not pull the average 181 # toward log10==0 (i.e. 1 ohm-m). Result should equal the positive-only 182 # interpolation. 183 with_bad = [ 184 self._s(-1.0, 0.0, 0.0, 100.0), 185 self._s(1.0, 0.0, 0.0, 0.0), 186 ] 187 positive_only = [self._s(-1.0, 0.0, 0.0, 100.0)] 188 result = idw(with_bad, [(0.0, 0.0, 0.0)], log10=True, radius=None) 189 expected = idw(positive_only, [(0.0, 0.0, 0.0)], log10=True, radius=None) 190 self.assertAlmostEqual(result[0], expected[0], places=6) 191 192 def test_multiple_exact_hits_are_averaged(self): 193 # Two samples coincide with the query; both must contribute (geometric 194 # mean in log space), not just whichever appears first. 195 samples = [ 196 self._s(0.0, 0.0, 0.0, 10.0), 197 self._s(0.0, 0.0, 0.0, 1000.0), 198 ] 199 result = idw(samples, [(0.0, 0.0, 0.0)], log10=True, radius=None) 200 self.assertAlmostEqual(result[0], math.sqrt(10.0 * 1000.0), places=4) 201 202 203 class DefaultRadiusTests(unittest.TestCase): 204 def _s(self, x, y): 205 return Sample(x, y, 0.0, 100.0, None) 206 207 def test_independent_of_sample_order(self): 208 locations = [(0.0, 0.0), (10.0, 0.0), (0.0, 10.0), (10.0, 10.0)] 209 forward = [self._s(x, y) for x, y in locations] 210 reversed_samples = list(reversed(forward)) 211 self.assertAlmostEqual( 212 _default_radius(forward), _default_radius(reversed_samples) 213 ) 214 215 def test_ignores_duplicate_depth_layers(self): 216 # Many depth layers at the same (x, y) must not collapse the radius to 0. 217 stacked = [Sample(0.0, 0.0, z, 100.0, None) for z in range(20)] 218 stacked += [Sample(30.0, 0.0, z, 100.0, None) for z in range(20)] 219 # Median nearest-neighbor spacing is 30; radius is 3x that. 220 self.assertAlmostEqual(_default_radius(stacked), 90.0) 221 222 223 class DefaultRadiusZTests(unittest.TestCase): 224 def test_stacked_column_median_gap(self): 225 # A single sounding stacked at 10 m spacing -> median gap 10 -> 1.5x. 226 column = [Sample(0.0, 0.0, z, 100.0, None) 227 for z in range(0, 100, 10)] 228 self.assertAlmostEqual(_default_radius_z(column), 15.0) 229 230 def test_returns_none_without_two_depths(self): 231 single = [Sample(0.0, 0.0, 0.0, 100.0, None)] 232 self.assertIsNone(_default_radius_z(single)) 233 234 235 class SamplesFromLayersTests(unittest.TestCase): 236 def _row(self, x, y, zmid, res): 237 return { 238 'X': x, 'Y': y, 'ZMid': zmid, 239 'Resistivity': res, 240 'Geometry': f'LINESTRING Z ({x} {y} {zmid + 1}, {x} {y} {zmid - 1})', 241 } 242 243 def test_basic_extraction(self): 244 rows = [self._row(1.0, 2.0, -5.0, 100.0)] 245 result = samples_from_layers(rows) 246 self.assertEqual(len(result), 1) 247 self.assertEqual(result[0].x, 1.0) 248 self.assertEqual(result[0].y, 2.0) 249 self.assertEqual(result[0].z, -5.0) 250 self.assertEqual(result[0].resistivity, 100.0) 251 self.assertIsNone(result[0].doi_elev) 252 253 def test_doi_join_by_xy(self): 254 rows = [self._row(1.0, 2.0, -5.0, 100.0)] 255 doi_by_xy = {(1.0, 2.0): -30.0} 256 result = samples_from_layers(rows, doi_by_xy) 257 self.assertEqual(result[0].doi_elev, -30.0) 258 259 def test_skips_nan_resistivity(self): 260 rows = [ 261 self._row(1.0, 2.0, -5.0, float('nan')), 262 self._row(3.0, 4.0, -5.0, 50.0), 263 ] 264 result = samples_from_layers(rows) 265 self.assertEqual(len(result), 1) 266 self.assertEqual(result[0].x, 3.0) 267 268 def test_skips_none_resistivity(self): 269 rows = [{'X': 1.0, 'Y': 2.0, 'ZMid': -5.0, 'Resistivity': None, 270 'Geometry': ''}] 271 result = samples_from_layers(rows) 272 self.assertEqual(len(result), 0) 273 274 275 class DoiByXYTests(unittest.TestCase): 276 def test_dict_rows(self): 277 doi_points = [{'X': 1.0, 'Y': 2.0, 'ZDOI': -30.0}] 278 result = doi_by_xy_from_points(doi_points) 279 self.assertEqual(result, {(1.0, 2.0): -30.0}) 280 281 def test_multiple_points(self): 282 doi_points = [ 283 {'X': 0.0, 'Y': 0.0, 'ZDOI': -10.0}, 284 {'X': 1.0, 'Y': 1.0, 'ZDOI': -20.0}, 285 ] 286 result = doi_by_xy_from_points(doi_points) 287 self.assertEqual(result[(0.0, 0.0)], -10.0) 288 self.assertEqual(result[(1.0, 1.0)], -20.0) 289 290 291 class InterpolateVoxelsTests(unittest.TestCase): 292 def _make_layer_rows(self, x, y, zmid, res, doi_elev=None): 293 return { 294 'X': x, 'Y': y, 'ZMid': zmid, 295 'Resistivity': res, 296 'Geometry': f'LINESTRING Z ({x} {y} {zmid + 1}, {x} {y} {zmid - 1})', 297 } 298 299 def test_basic_output_keys(self): 300 rows = [self._make_layer_rows(0.0, 0.0, 0.0, 100.0)] 301 g = GridSpec(0.0, 0.0, 0.0, 10.0, 10.0, 10.0, 1, 1, 1) 302 voxels = interpolate_voxels(rows, g) 303 self.assertEqual(len(voxels), 1) 304 v = voxels[0] 305 for key in ('X', 'Y', 'Z', 'Resistivity', 'Color', 'CellX', 'CellY', 306 'CellZ', 'NumSamples', 'Geometry'): 307 self.assertIn(key, v) 308 309 def test_color_set_from_ramp(self): 310 from tem_loader.core import resistivity_color as rc 311 rows = [self._make_layer_rows(0.0, 0.0, 0.0, 31.0)] 312 g = GridSpec(0.0, 0.0, 0.0, 10.0, 10.0, 10.0, 1, 1, 1) 313 voxels = interpolate_voxels(rows, g) 314 self.assertEqual(voxels[0]['Color'], rc(voxels[0]['Resistivity'])) 315 316 def test_geometry_is_point_z_wkt(self): 317 rows = [self._make_layer_rows(0.0, 0.0, 0.0, 50.0)] 318 g = GridSpec(0.0, 0.0, 0.0, 10.0, 10.0, 10.0, 1, 1, 1) 319 voxels = interpolate_voxels(rows, g) 320 self.assertTrue(voxels[0]['Geometry'].startswith('POINT Z')) 321 322 def test_num_samples_reported(self): 323 rows = [ 324 self._make_layer_rows(0.0, 0.0, 0.0, 50.0), 325 self._make_layer_rows(1.0, 1.0, 0.0, 100.0), 326 ] 327 # Cell center is at (1, 1, 1), close to both samples. 328 g = GridSpec(0.0, 0.0, 0.0, 2.0, 2.0, 2.0, 1, 1, 1) 329 voxels = interpolate_voxels(rows, g) 330 self.assertTrue(len(voxels) > 0, 'no voxels produced') 331 self.assertGreaterEqual(voxels[0]['NumSamples'], 1) 332 333 def test_per_sample_doi_gating(self): 334 # Sounding A (0,0) has DOI at z=-60; sounding B (5,0) has DOI at z=-80. 335 # Cells: zmin=-90, dz=30, nz=3 → centers at z=-75, -45, -15. 336 # CellZ=0 (z=-75): A gated (−75 < −60), B usable (−75 > −80). B-only. 337 # CellZ=1 (z=-45): both usable. 338 # CellZ=2 (z=-15): both usable. 339 rows = [ 340 self._make_layer_rows(0.0, 0.0, -10.0, 50.0), 341 self._make_layer_rows(5.0, 0.0, -10.0, 200.0), 342 ] 343 doi_dict = [ 344 {'X': 0.0, 'Y': 0.0, 'ZDOI': -60.0}, 345 {'X': 5.0, 'Y': 0.0, 'ZDOI': -80.0}, 346 ] 347 g = GridSpec(-5.0, -5.0, -90.0, 20.0, 20.0, 30.0, 1, 1, 3) 348 349 voxels_by_z = {v['CellZ']: v for v in interpolate_voxels( 350 rows, g, doi_points=doi_dict, clip_below_doi=True, min_neighbors=1 351 )} 352 # CellZ=0 (z=-75): only B should contribute 353 b_only = voxels_by_z.get(0) 354 self.assertIsNotNone(b_only, 'expected a voxel in B-only zone at z=-75') 355 self.assertEqual(b_only['NumSamples'], 1) 356 # CellZ=1 (z=-45): both should contribute 357 both = voxels_by_z.get(1) 358 self.assertIsNotNone(both, 'expected a voxel in both-usable zone at z=-45') 359 self.assertGreaterEqual(both['NumSamples'], 2) 360 361 def test_thin_layer_stays_sharp_with_radius_z(self): 362 # A single sounding with a sharp resistivity contrast at two depths. 363 # A tight vertical search radius must keep each voxel pinned to its own 364 # depth instead of bleeding the two values together across CellZ. 365 rows = [ 366 self._make_layer_rows(0.0, 0.0, 0.0, 10.0), 367 self._make_layer_rows(0.0, 0.0, -50.0, 1000.0), 368 ] 369 g = GridSpec(-5.0, -5.0, -60.0, 10.0, 10.0, 10.0, 1, 1, 6) 370 voxels_by_z = {v['CellZ']: v for v in interpolate_voxels( 371 rows, g, radius=20.0, radius_z=10.0 372 )} 373 top = voxels_by_z.get(5) # center z = -5, nearest the 10 ohm-m sample 374 bottom = voxels_by_z.get(0) # center z = -55, nearest the 1000 ohm-m 375 self.assertIsNotNone(top) 376 self.assertIsNotNone(bottom) 377 self.assertAlmostEqual(top['Resistivity'], 10.0, places=6) 378 self.assertAlmostEqual(bottom['Resistivity'], 1000.0, places=6) 379 self.assertEqual(top['NumSamples'], 1) 380 self.assertEqual(bottom['NumSamples'], 1) 381 382 def test_safety_cap_raises(self): 383 bounds = (0.0, 200.0, 0.0, 200.0, 0.0, 200.0) 384 g = GridSpec.from_cell_size(bounds, 1.0, 1.0, 1.0) 385 rows = [self._make_layer_rows(100.0, 100.0, 100.0, 50.0)] 386 with self.assertRaises(ValueError): 387 interpolate_voxels(rows, g) 388 389 def test_progress_callback(self): 390 rows = [self._make_layer_rows(0.0, 0.0, 0.0, 50.0)] 391 g = GridSpec(0.0, 0.0, 0.0, 10.0, 10.0, 10.0, 2, 2, 2) 392 calls = [] 393 interpolate_voxels(rows, g, progress=lambda done, total: calls.append((done, total))) 394 self.assertTrue(all(total == g.cell_count() for _, total in calls)) 395 self.assertEqual(calls[-1][0], g.cell_count()) 396 397 def test_progress_truthy_return_cancels(self): 398 rows = [self._make_layer_rows(0.0, 0.0, 0.0, 50.0)] 399 g = GridSpec(0.0, 0.0, 0.0, 10.0, 10.0, 10.0, 4, 4, 4) 400 calls = [] 401 402 def cancel_after_first(done, total): 403 calls.append(done) 404 return done >= 1 405 406 interpolate_voxels(rows, g, progress=cancel_after_first) 407 self.assertEqual(calls, [1]) 408 409 def test_matches_idw_values(self): 410 # interpolate_voxels and idw share one core, so a voxel's resistivity 411 # must equal idw evaluated at the same cell center. 412 rows = [ 413 self._make_layer_rows(0.0, 0.0, 0.0, 50.0), 414 self._make_layer_rows(20.0, 0.0, -5.0, 500.0), 415 ] 416 g = GridSpec(-10.0, -10.0, -10.0, 10.0, 10.0, 10.0, 3, 1, 2) 417 voxels = interpolate_voxels(rows, g, power=2.0, log10=True) 418 samples = samples_from_layers(rows) 419 for v in voxels: 420 expected = idw(samples, [(v['X'], v['Y'], v['Z'])], 421 power=2.0, log10=True)[0] 422 self.assertAlmostEqual(v['Resistivity'], expected, places=6) 423 424 425 if __name__ == '__main__': 426 unittest.main()