test_mean_resistivity.py (10232B)
1 import math 2 import unittest 3 from pathlib import Path 4 5 import numpy as np 6 7 from tem_loader.core import process_xyz 8 from tem_loader.mean_resistivity import ( 9 MapGrid, 10 elevation_intervals, 11 idw_grid, 12 interval_mean, 13 map_grid, 14 mean_resistivity_maps, 15 soundings_from_layers, 16 ) 17 18 FIXTURE_DIR = Path(__file__).parent / 'data' 19 20 21 def _layer(x, y, ztop, zbottom, res): 22 return {'X': x, 'Y': y, 'ZTop': ztop, 'ZBottom': zbottom, 23 'Resistivity': res} 24 25 26 class ElevationIntervalTests(unittest.TestCase): 27 def test_whole_multiples(self): 28 self.assertEqual( 29 elevation_intervals(20.0, -10.0, 10.0), 30 [(20.0, 10.0), (10.0, 0.0), (0.0, -10.0)], 31 ) 32 33 def test_last_interval_truncated(self): 34 self.assertEqual( 35 elevation_intervals(20.0, 5.0, 10.0), 36 [(20.0, 10.0), (10.0, 5.0)], 37 ) 38 39 def test_rejects_invalid_input(self): 40 with self.assertRaises(ValueError): 41 elevation_intervals(20.0, 0.0, 0.0) 42 with self.assertRaises(ValueError): 43 elevation_intervals(0.0, 0.0, 10.0) 44 with self.assertRaises(ValueError): 45 elevation_intervals(-10.0, 0.0, 10.0) 46 47 48 class MapGridTests(unittest.TestCase): 49 def test_extent_snapped_to_cell_size(self): 50 grid = map_grid([1012.0, 1188.0], [2049.0, 2101.0], 50.0) 51 self.assertEqual(grid, MapGrid(1000.0, 2150.0, 50.0, 4, 3)) 52 53 def test_covers_all_points(self): 54 xs = [3.3, 97.1, 51.0] 55 ys = [-12.0, 40.5, 7.7] 56 grid = map_grid(xs, ys, 10.0) 57 self.assertLessEqual(grid.xmin, min(xs)) 58 self.assertGreaterEqual(grid.xmin + grid.nx * grid.cell_size, max(xs)) 59 self.assertGreaterEqual(grid.ymax, max(ys)) 60 self.assertLessEqual(grid.ymax - grid.ny * grid.cell_size, min(ys)) 61 62 def test_single_point(self): 63 grid = map_grid([100.0], [200.0], 25.0) 64 self.assertEqual((grid.nx, grid.ny), (1, 1)) 65 66 def test_rejects_invalid_input(self): 67 with self.assertRaises(ValueError): 68 map_grid([0.0], [0.0], 0.0) 69 with self.assertRaises(ValueError): 70 map_grid([], [], 10.0) 71 with self.assertRaises(ValueError): 72 map_grid([0.0, 1e6], [0.0, 1e6], 1.0) 73 74 75 class SoundingsFromLayersTests(unittest.TestCase): 76 def test_groups_by_xy(self): 77 rows = [ 78 _layer(0.0, 0.0, 10.0, 5.0, 100.0), 79 _layer(0.0, 0.0, 5.0, 0.0, 10.0), 80 _layer(50.0, 0.0, 10.0, 0.0, 30.0), 81 ] 82 soundings = sorted(soundings_from_layers(rows)) 83 self.assertEqual(len(soundings), 2) 84 self.assertEqual(soundings[0].res.tolist(), [100.0, 10.0]) 85 self.assertEqual(soundings[1].ztop.tolist(), [10.0]) 86 87 def test_skips_invalid_resistivity(self): 88 rows = [ 89 _layer(0.0, 0.0, 10.0, 5.0, None), 90 _layer(0.0, 0.0, 10.0, 5.0, math.nan), 91 _layer(0.0, 0.0, 10.0, 5.0, 0.0), 92 _layer(0.0, 0.0, 10.0, 5.0, -5.0), 93 ] 94 self.assertEqual(soundings_from_layers(rows), []) 95 96 def test_clips_to_doi(self): 97 rows = [ 98 _layer(0.0, 0.0, 10.0, 0.0, 100.0), 99 _layer(0.0, 0.0, 0.0, -20.0, 10.0), 100 ] 101 (sounding,) = soundings_from_layers(rows, {(0.0, 0.0): 4.0}) 102 self.assertEqual(sounding.ztop.tolist(), [10.0]) 103 self.assertEqual(sounding.zbottom.tolist(), [4.0]) 104 105 106 class IntervalMeanTests(unittest.TestCase): 107 def _mean(self, layers, top, bottom, min_coverage=0.5): 108 arr = np.asarray(layers, dtype=float) 109 return interval_mean(arr[:, 0], arr[:, 1], arr[:, 2], top, bottom, 110 min_coverage) 111 112 def test_single_covering_layer(self): 113 value = self._mean([(50.0, -50.0, 40.0)], 10.0, 0.0) 114 self.assertAlmostEqual(10.0 ** value, 40.0) 115 116 def test_equal_thickness_layers_give_geometric_mean(self): 117 value = self._mean([(10.0, 5.0, 10.0), (5.0, 0.0, 1000.0)], 10.0, 0.0) 118 self.assertAlmostEqual(10.0 ** value, 100.0) 119 120 def test_partial_overlap_weighted_by_thickness(self): 121 layers = [(20.0, 8.0, 10.0), (8.0, -20.0, 1000.0)] 122 value = self._mean(layers, 10.0, 0.0) 123 self.assertAlmostEqual(value, (2.0 * 1.0 + 8.0 * 3.0) / 10.0) 124 125 def test_low_coverage_returns_none(self): 126 self.assertIsNone(self._mean([(10.0, 6.0, 10.0)], 10.0, 0.0)) 127 self.assertIsNotNone(self._mean([(10.0, 6.0, 10.0)], 10.0, 0.0, 128 min_coverage=0.4)) 129 130 def test_no_overlap_returns_none(self): 131 self.assertIsNone(self._mean([(30.0, 20.0, 10.0)], 10.0, 0.0, 132 min_coverage=0.0)) 133 134 135 class IdwGridTests(unittest.TestCase): 136 def setUp(self): 137 self.grid = MapGrid(0.0, 10.0, 10.0, 3, 1) 138 139 def test_node_at_point_returns_its_value(self): 140 out = idw_grid([5.0, 25.0], [5.0, 5.0], [1.0, 3.0], self.grid, 141 radius=100.0) 142 self.assertAlmostEqual(out[0, 0], 1.0) 143 self.assertAlmostEqual(out[0, 2], 3.0) 144 145 def test_midpoint_is_mean(self): 146 out = idw_grid([5.0, 25.0], [5.0, 5.0], [1.0, 3.0], self.grid, 147 radius=100.0) 148 self.assertAlmostEqual(out[0, 1], 2.0) 149 150 def test_nodes_beyond_radius_are_blank(self): 151 out = idw_grid([5.0], [5.0], [1.0], self.grid, radius=12.0) 152 self.assertAlmostEqual(out[0, 0], 1.0) 153 self.assertAlmostEqual(out[0, 1], 1.0) 154 self.assertTrue(math.isnan(out[0, 2])) 155 156 def test_min_neighbors_enforced(self): 157 out = idw_grid([5.0, 25.0], [5.0, 5.0], [1.0, 3.0], self.grid, 158 radius=12.0, min_neighbors=2) 159 self.assertTrue(math.isnan(out[0, 0])) 160 self.assertAlmostEqual(out[0, 1], 2.0) 161 self.assertTrue(math.isnan(out[0, 2])) 162 163 def test_unlimited_radius_uses_all_points(self): 164 out = idw_grid([5.0], [5.0], [1.0], self.grid) 165 self.assertTrue(np.all(out == 1.0)) 166 167 def test_rows_are_north_up(self): 168 grid = MapGrid(0.0, 20.0, 10.0, 1, 2) 169 out = idw_grid([5.0, 5.0], [15.0, 5.0], [1.0, 2.0], grid, 170 radius=1.0) 171 self.assertEqual(out[:, 0].tolist(), [1.0, 2.0]) 172 173 def test_large_grid_spans_tiles(self): 174 grid = MapGrid(0.0, 1000.0, 10.0, 100, 100) 175 out = idw_grid([5.0, 995.0], [995.0, 5.0], [1.0, 2.0], grid, 176 radius=15.0) 177 self.assertAlmostEqual(out[0, 0], 1.0) 178 self.assertAlmostEqual(out[99, 99], 2.0) 179 self.assertEqual(int(np.isfinite(out).sum()), 2 * 4) 180 181 182 class MeanResistivityMapsTests(unittest.TestCase): 183 def setUp(self): 184 self.rows = [] 185 for x, res_top, res_bottom in ((0.0, 10.0, 100.0), 186 (100.0, 1000.0, 100.0)): 187 self.rows.append(_layer(x, 0.0, 10.0, 0.0, res_top)) 188 self.rows.append(_layer(x, 0.0, 0.0, -10.0, res_bottom)) 189 190 def test_maps_per_interval(self): 191 grid, maps = mean_resistivity_maps( 192 self.rows, [(10.0, 0.0), (0.0, -10.0)], 50.0, radius=500.0, 193 ) 194 self.assertEqual((grid.nx, grid.ny), (2, 1)) 195 self.assertEqual([(m[0], m[1]) for m in maps], 196 [(10.0, 0.0), (0.0, -10.0)]) 197 top_map = maps[0][2] 198 self.assertEqual(top_map.shape, (1, 2)) 199 self.assertTrue(np.all(top_map > 10.0)) 200 self.assertTrue(np.all(top_map < 1000.0)) 201 np.testing.assert_allclose(maps[1][2], 100.0, rtol=1e-6) 202 self.assertEqual(maps[0][3], 2) 203 204 def test_intervals_without_data_are_omitted(self): 205 _, maps = mean_resistivity_maps( 206 self.rows, [(50.0, 40.0), (10.0, 0.0), (-50.0, -60.0)], 50.0, 207 ) 208 self.assertEqual([(m[0], m[1]) for m in maps], [(10.0, 0.0)]) 209 210 def test_doi_points_clip_layers(self): 211 doi = [{'X': 0.0, 'Y': 0.0, 'ZDOI': 0.0}, 212 {'X': 100.0, 'Y': 0.0, 'ZDOI': 0.0}] 213 _, maps = mean_resistivity_maps( 214 self.rows, [(10.0, 0.0), (0.0, -10.0)], 50.0, doi_points=doi, 215 ) 216 self.assertEqual([(m[0], m[1]) for m in maps], [(10.0, 0.0)]) 217 218 def test_progress_can_cancel(self): 219 calls = [] 220 221 def progress(done, total): 222 calls.append((done, total)) 223 return True 224 225 _, maps = mean_resistivity_maps( 226 self.rows, [(10.0, 0.0), (0.0, -10.0)], 50.0, progress=progress, 227 ) 228 self.assertEqual(calls, [(1, 2)]) 229 self.assertEqual(len(maps), 1) 230 231 def test_rejects_invalid_parameters(self): 232 with self.assertRaises(ValueError): 233 mean_resistivity_maps(self.rows, [(10.0, 0.0)], 50.0, 234 min_coverage=1.5) 235 with self.assertRaises(ValueError): 236 mean_resistivity_maps(self.rows, [(10.0, 0.0)], 50.0, radius=0.0) 237 with self.assertRaises(ValueError): 238 mean_resistivity_maps([], [(10.0, 0.0)], 50.0) 239 240 def test_fixtures(self): 241 for name in ('stem_80x80_workbench_2024_2_0_0.xyz', 242 'stem_temimager_4_0_4_6.xyz'): 243 with self.subTest(name=name): 244 _, doi_points, layers = process_xyz(FIXTURE_DIR / name) 245 tops = [row['ZTop'] for row in layers] 246 bottoms = [row['ZBottom'] for row in layers] 247 intervals = elevation_intervals( 248 math.ceil(max(tops) / 10.0) * 10.0, 249 math.floor(min(bottoms) / 10.0) * 10.0, 250 10.0, 251 ) 252 _, maps = mean_resistivity_maps( 253 layers, intervals, 100.0, 254 doi_points=doi_points or None, 255 ) 256 self.assertTrue(maps) 257 self.assertLess(len(maps), len(intervals) + 1) 258 res = [row['Resistivity'] for row in layers 259 if row['Resistivity'] > 0] 260 for top, bottom, array, count in maps: 261 finite = array[np.isfinite(array)] 262 self.assertGreater(count, 0) 263 self.assertGreaterEqual(finite.min(), min(res) * 0.999) 264 self.assertLessEqual(finite.max(), max(res) * 1.001) 265 266 267 if __name__ == '__main__': 268 unittest.main()