qgis-tem-loader

commit 99c3ac3a0f150a3271b3b16e357d3efcd8048b80
parent 200121e8731ef36cb89dccfd6a7f04176bcc96b0
Author: Anders Damsgaard <anders@adamsgaard.dk>
Date:   Thu, 22 Oct 2026 09:05:32 +0000

feat(maps): write mean resistivity maps as styled GeoTIFFs

Diffstat:
A
tem_loader/geotiff.py
 | 
27
+++++++++++++++++++++++++++
A
tem_loader/styles/mean_resistivity.qml
 | 
52
++++++++++++++++++++++++++++++++++++++++++++++++++++
M
test/test_core.py
 | 
13
+++++++++++++
A
test/test_geotiff.py
 | 
55
+++++++++++++++++++++++++++++++++++++++++++++++++++++++

4 files changed, 147 insertions(+), 0 deletions(-)
diff --git a/tem_loader/geotiff.py b/tem_loader/geotiff.py
@@ -0,0 +1,27 @@
+import numpy as np
+from osgeo import gdal
+
+NODATA = -9999.0
+
+
+def write_geotiff(path, grid, array, crs_wkt):
+    """Write a north-up single-band Float32 GeoTIFF; NaN cells become NODATA."""
+    data = np.where(np.isfinite(array), array, NODATA).astype(np.float32)
+    with gdal.ExceptionMgr(useExceptions=True):
+        try:
+            driver = gdal.GetDriverByName('GTiff')
+            dataset = driver.Create(
+                str(path), grid.nx, grid.ny, 1, gdal.GDT_Float32,
+                options=['COMPRESS=DEFLATE', 'PREDICTOR=3'],
+            )
+            dataset.SetGeoTransform((grid.xmin, grid.cell_size, 0.0,
+                                     grid.ymax, 0.0, -grid.cell_size))
+            if crs_wkt:
+                dataset.SetProjection(crs_wkt)
+            band = dataset.GetRasterBand(1)
+            band.SetNoDataValue(NODATA)
+            band.WriteArray(data)
+            band.FlushCache()
+            dataset = None
+        except RuntimeError as exc:
+            raise ValueError(f'failed to write GeoTIFF {path}: {exc}') from exc
diff --git a/tem_loader/styles/mean_resistivity.qml b/tem_loader/styles/mean_resistivity.qml
@@ -0,0 +1,52 @@
+<!DOCTYPE qgis PUBLIC 'http://mrcc.com/qgis.dtd' 'SYSTEM'>
+<qgis hasScaleBasedVisibilityFlag="0" maxScale="0" minScale="1e+08" styleCategories="AllStyleCategories" version="3.40.0-Bratislava">
+  <pipe>
+    <provider>
+      <resampling enabled="false" maxOversampling="2" zoomedInResamplingMethod="nearestNeighbour" zoomedOutResamplingMethod="nearestNeighbour"/>
+    </provider>
+    <rasterrenderer alphaBand="-1" band="1" classificationMax="inf" classificationMin="0" nodataColor="" opacity="1" type="singlebandpseudocolor">
+      <rasterTransparency/>
+      <minMaxOrigin>
+        <limits>None</limits>
+        <extent>WholeRaster</extent>
+        <statAccuracy>Estimated</statAccuracy>
+        <cumulativeCutLower>0.02</cumulativeCutLower>
+        <cumulativeCutUpper>0.98</cumulativeCutUpper>
+        <stdDevFactor>2</stdDevFactor>
+      </minMaxOrigin>
+      <rastershader>
+        <colorrampshader classificationMode="1" clip="0" colorRampType="DISCRETE" labelPrecision="1" maximumValue="inf" minimumValue="0">
+          <item alpha="255" color="#000091" label="&lt; 0.1" value="0.1"/>
+          <item alpha="255" color="#000091" label="0.1 - 1" value="1"/>
+          <item alpha="255" color="#0000b4" label="1 - 2" value="2"/>
+          <item alpha="255" color="#0032dc" label="2 - 4" value="4"/>
+          <item alpha="255" color="#005af5" label="4 - 7" value="7"/>
+          <item alpha="255" color="#008cff" label="7 - 10" value="10"/>
+          <item alpha="255" color="#00beff" label="10 - 15" value="15"/>
+          <item alpha="255" color="#00dcff" label="15 - 20" value="20"/>
+          <item alpha="255" color="#00ffff" label="20 - 25" value="25"/>
+          <item alpha="255" color="#00ff96" label="25 - 30" value="30"/>
+          <item alpha="255" color="#00ff00" label="30 - 35" value="35"/>
+          <item alpha="255" color="#96ff00" label="35 - 40" value="40"/>
+          <item alpha="255" color="#d2ff00" label="40 - 50" value="50"/>
+          <item alpha="255" color="#ffff00" label="50 - 60" value="60"/>
+          <item alpha="255" color="#ffb500" label="60 - 75" value="75"/>
+          <item alpha="255" color="#ff7300" label="75 - 90" value="90"/>
+          <item alpha="255" color="#ff0000" label="90 - 120" value="120"/>
+          <item alpha="255" color="#ff1c8d" label="120 - 150" value="150"/>
+          <item alpha="255" color="#ff6aff" label="150 - 200" value="200"/>
+          <item alpha="255" color="#f200f2" label="200 - 250" value="250"/>
+          <item alpha="255" color="#ca00ca" label="250 - 300" value="300"/>
+          <item alpha="255" color="#a600a6" label="300 - 400" value="400"/>
+          <item alpha="255" color="#800080" label="400 - 600" value="600"/>
+          <item alpha="255" color="#750075" label="600 - 1600" value="1600"/>
+          <item alpha="255" color="#540054" label="&gt; 1600" value="inf"/>
+        </colorrampshader>
+      </rastershader>
+    </rasterrenderer>
+    <brightnesscontrast brightness="0" contrast="0" gamma="1"/>
+    <huesaturation colorizeBlue="128" colorizeGreen="128" colorizeOn="0" colorizeRed="255" colorizeStrength="100" grayscaleMode="0" invertColors="0" saturation="0"/>
+    <rasterresampler maxOversampling="2"/>
+  </pipe>
+  <blendMode>0</blendMode>
+</qgis>
diff --git a/test/test_core.py b/test/test_core.py
@@ -12,6 +12,7 @@ import xml.etree.ElementTree as ET
 from tem_loader.core import (
     ABOVE_DOI_OPACITY,
     BELOW_DOI_OPACITY,
+    RESISTIVITY_COLOR_RAMP,
     detect_source_epsg,
     layer_opacity,
     process_xyz,
@@ -396,6 +397,18 @@ class ProcessXYZTests(unittest.TestCase):
         self.assertEqual(active.attrib["value"], "true")
         self.assertEqual(expr.attrib["value"], '"Color"')
 
+    def test_mean_resistivity_style_mirrors_color_ramp(self):
+        tree = ET.parse(STYLES_DIR / "mean_resistivity.qml")
+        renderer = tree.getroot().find(".//rasterrenderer")
+        self.assertEqual(renderer.attrib["type"], "singlebandpseudocolor")
+        shader = renderer.find("./rastershader/colorrampshader")
+        self.assertEqual(shader.attrib["colorRampType"], "DISCRETE")
+        self.assertEqual(
+            [(float(item.attrib["value"]), item.attrib["color"])
+             for item in shader.findall("./item")],
+            list(RESISTIVITY_COLOR_RAMP),
+        )
+
     def test_detect_source_epsg_for_sci_fixture(self):
         path = FIXTURE_DIR / "sci_workbench_2026_1.xyz"
         self.assertEqual(detect_source_epsg(path), "EPSG:25832")
diff --git a/test/test_geotiff.py b/test/test_geotiff.py
@@ -0,0 +1,55 @@
+import math
+import unittest
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+
+from tem_loader.mean_resistivity import MapGrid
+
+try:
+    from osgeo import gdal, osr
+except ImportError:
+    gdal = None
+
+if gdal is not None:
+    osr.UseExceptions()
+    from tem_loader.geotiff import NODATA, write_geotiff
+
+
+@unittest.skipIf(gdal is None, 'GDAL Python bindings not available')
+class WriteGeotiffTests(unittest.TestCase):
+    def test_round_trip(self):
+        grid = MapGrid(500000.0, 6200000.0, 50.0, 3, 2)
+        array = np.array([[10.0, 20.0, math.nan], [30.0, 40.0, 50.0]])
+        srs = osr.SpatialReference()
+        srs.ImportFromEPSG(25832)
+
+        with TemporaryDirectory() as tmp:
+            path = Path(tmp) / 'map.tif'
+            write_geotiff(path, grid, array, srs.ExportToWkt())
+            dataset = gdal.Open(str(path))
+            band = dataset.GetRasterBand(1)
+
+            self.assertEqual((dataset.RasterXSize, dataset.RasterYSize),
+                             (3, 2))
+            self.assertEqual(dataset.GetGeoTransform(),
+                             (500000.0, 50.0, 0.0, 6200000.0, 0.0, -50.0))
+            self.assertEqual(band.DataType, gdal.GDT_Float32)
+            self.assertEqual(band.GetNoDataValue(), NODATA)
+            self.assertEqual(band.ReadAsArray().tolist(),
+                             [[10.0, 20.0, NODATA], [30.0, 40.0, 50.0]])
+            written = osr.SpatialReference(wkt=dataset.GetProjection())
+            self.assertEqual(written.GetAuthorityCode(None), '25832')
+            dataset = None
+
+    def test_unwritable_path_raises_value_error(self):
+        grid = MapGrid(0.0, 10.0, 10.0, 1, 1)
+        with TemporaryDirectory() as tmp:
+            path = Path(tmp) / 'missing' / 'map.tif'
+            with self.assertRaises(ValueError):
+                write_geotiff(path, grid, np.ones((1, 1)), '')
+
+
+if __name__ == '__main__':
+    unittest.main()