Source code: 3D-print a geological ground model: one watertight STL per stratum

"""3D-print a geological ground model: one watertight STL per stratum, plus a colour 3MF.

Part 7 of the TechnicalGIS series "AGS data in GIS". Reads the clamped stratum surfaces that Part 5
interpolated from the A417 Missing Link boreholes, cuts a window over the Cotswold escarpment at
Crickley Hill, and turns it into a multi-material block: superficial deposits, Great Oolite, Inferior
Oolite, Bridport Sand and Lias, each a separate solid that fits the others exactly. Units too thin to
print are handed to the unit below, every part is checked for watertightness before it is written,
and the previews are rendered from the STL files read back from disk.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/3d-print-geological-ground-model/

Data: A417 Missing Link ground investigation (BGS deposited data items 173984 and 173975; client
Highways England), from the BGS AGS download service under the Open Government Licence v3.0. Contains British Geological Survey materials © UKRI 2026.
Environment Agency LIDAR Composite DTM 1 m via the Defra WCS (Open Government Licence v3).
"""
import argparse
import json
import sys
import time
import zipfile
import xml.etree.ElementTree as ET
from importlib.metadata import version
from pathlib import Path

import numpy as np
import rasterio
from rasterio import features
from rasterio.windows import Window, from_bounds
from scipy import ndimage
import shapely
from shapely.geometry import shape
from shapely.geometry.polygon import orient
from shapely.ops import unary_union
import manifold3d as mf
import pyvista as pv
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch, PathPatch
from matplotlib.path import Path as MplPath

plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9,
                     "figure.facecolor": "white", "axes.facecolor": "white"})
pv.OFF_SCREEN = True
T0 = time.perf_counter()

MODEL = Path("../geological-ground-model-surfaces-python/output/model")   # written by Part 5
OUT = Path("output")
WINDOW = (392_400, 215_000, 394_300, 216_400)    # BNG window over the escarpment at Crickley Hill
LONG_MM = 200.0          # long side of the print: 198 x 145 mm fits a 220 mm build plate
VE = 3.0                 # vertical exaggeration, as in the 3D views of Parts 3 and 5
BASE_MM = 3.0            # plinth under the model base, printed with the bottom unit
MIN_MM = 0.6             # thinnest layer worth printing (about three 0.2 mm layers)
PLA_G_CM3 = 1.24
EXPLODE_MM = 14.0        # vertical gap between parts in the exploded view


def section(title):
    print(f"\n=== {title} ===")


# ------------------------------------------------------------------ mesh helpers (as in the multi-colour terrain tutorial)
def grid_triangles(r, c):
    idx = np.arange(r * c).reshape(r, c)
    a, b = idx[:-1, :-1].ravel(), idx[:-1, 1:].ravel()
    d, e = idx[1:, :-1].ravel(), idx[1:, 1:].ravel()
    return np.concatenate([np.column_stack([a, d, b]), np.column_stack([b, d, e])])


def heightfield_solid(Z, cell_mm):
    rows, cols = Z.shape
    X, Y = np.meshgrid(np.arange(cols) * cell_mm, (rows - 1 - np.arange(rows)) * cell_mm)
    top = np.column_stack([X.ravel(), Y.ravel(), Z.ravel()])
    bot = top.copy()
    bot[:, 2] = 0.0
    n = rows * cols
    tri_top = grid_triangles(rows, cols)
    idx = np.arange(n).reshape(rows, cols)
    ring = np.concatenate([idx[0, :], idx[1:, -1], idx[-1, -2::-1], idx[-2:0:-1, 0]])
    p, q = ring, np.roll(ring, -1)
    walls = np.concatenate([np.column_stack([p, q, p + n]), np.column_stack([q, q + n, p + n])])
    return np.vstack([top, bot]), np.vstack([tri_top, tri_top[:, ::-1] + n, walls])


def to_manifold(verts, tris):
    m = mf.Manifold(mf.Mesh(vert_properties=np.ascontiguousarray(verts, dtype=np.float32),
                            tri_verts=np.ascontiguousarray(tris, dtype=np.uint32)))
    assert m.status() == mf.Error.NoError, m.status()
    return m


def from_manifold(m):
    mesh = m.to_mesh()
    return np.asarray(mesh.vert_properties)[:, :3].astype("float64"), np.asarray(mesh.tri_verts).astype("int64")


def check_mesh(verts, tris):
    """Every edge shared by exactly two triangles, used once in each direction, and a positive volume."""
    e = np.concatenate([tris[:, [0, 1]], tris[:, [1, 2]], tris[:, [2, 0]]])
    n = len(verts)
    und = np.unique(np.sort(e, axis=1) @ np.array([n, 1]), return_counts=True)[1]
    drc = np.unique(e @ np.array([n, 1]), return_counts=True)[1]
    t = verts[tris]
    vol = np.einsum("ij,ij->i", t[:, 0], np.cross(t[:, 1], t[:, 2])).sum() / 6
    return bool((und == 2).all()), bool((drc == 1).all()), vol


def drop_slivers(m, min_mm3=0.05):
    """Booleans leave zero-volume slivers where two surfaces touch; keep only real pieces."""
    pieces = [p for p in m.decompose() if p.volume() >= min_mm3]
    return mf.Manifold.batch_boolean(pieces, mf.OpType.Add) if pieces else mf.Manifold(), len(pieces)


def write_stl(path, verts, tris, header):
    t = verts[tris]
    nrm = np.cross(t[:, 1] - t[:, 0], t[:, 2] - t[:, 0])
    ln = np.linalg.norm(nrm, axis=1, keepdims=True)
    rec = np.zeros(len(tris), dtype=[("normal", "<f4", 3), ("v", "<f4", (3, 3)), ("attr", "<u2")])
    rec["normal"] = np.divide(nrm, ln, out=np.zeros_like(nrm), where=ln > 0)
    rec["v"] = t
    with open(path, "wb") as f:
        f.write(header.encode("ascii")[:80].ljust(80, b" "))
        f.write(np.uint32(len(rec)).tobytes())
        f.write(rec.tobytes())


def write_3mf(path, parts):
    """Minimal 3MF: one mesh object per part, each with its own display colour."""
    ns = "http://schemas.microsoft.com/3dmanufacturing/core/2015/02"
    xml = [f'<?xml version="1.0" encoding="UTF-8"?>\n<model unit="millimeter" xml:lang="en-GB" xmlns="{ns}">',
           "<resources>", '<basematerials id="1">']
    xml += [f'<base name="{name}" displaycolor="{colour.upper()}FF"/>' for name, _, _, colour in parts]
    xml.append("</basematerials>")
    for i, (name, v, t, _) in enumerate(parts):
        xml.append(f'<object id="{i + 2}" name="{name}" type="model" pid="1" pindex="{i}"><mesh><vertices>')
        xml.append("".join(f'<vertex x="{x:.4f}" y="{y:.4f}" z="{z:.4f}"/>' for x, y, z in v))
        xml.append("</vertices><triangles>")
        xml.append("".join(f'<triangle v1="{a}" v2="{b}" v3="{c}"/>' for a, b, c in t))
        xml.append("</triangles></mesh></object>")
    xml += ["</resources>", "<build>"] + [f'<item objectid="{i + 2}"/>' for i in range(len(parts))]
    xml += ["</build>", "</model>"]
    with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as z:
        z.writestr("[Content_Types].xml",
                   '<?xml version="1.0" encoding="UTF-8"?><Types xmlns="http://schemas.openxmlformats.org/'
                   'package/2006/content-types"><Default Extension="rels" ContentType="application/'
                   'vnd.openxmlformats-package.relationships+xml"/><Default Extension="model" '
                   'ContentType="application/vnd.ms-package.3dmanufacturing-3dmodel+xml"/></Types>')
        z.writestr("_rels/.rels",
                   '<?xml version="1.0" encoding="UTF-8"?><Relationships xmlns="http://schemas.openxmlformats.org/'
                   'package/2006/relationships"><Relationship Target="/3D/3dmodel.model" Id="rel0" '
                   'Type="http://schemas.microsoft.com/3dmanufacturing/2013/01/3dmodel"/></Relationships>')
        z.writestr("3D/3dmodel.model", "\n".join(xml))


parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--model", type=Path, default=MODEL, help="folder with Part 5's surfaces and model.json")
args = parser.parse_args()
if not (args.model / "model.json").exists():
    sys.exit("Run Part 5 (geological-ground-model-surfaces-python) first – this script reads its output/model/ surfaces.")
OUT.mkdir(exist_ok=True)

print(f"numpy {np.__version__} | rasterio {rasterio.__version__} | shapely {version('shapely')} | "
      f"manifold3d {version('manifold3d')} | pyvista {pv.__version__}")

# ------------------------------------------------------------------ Step 1: surfaces in the print window
section("Step 1: read the stratum surfaces in the print window")
meta = json.loads((args.model / "model.json").read_text(encoding="utf-8"))
units, base_m = meta["units"], float(meta["model_base_mOD"])
for upper, lower in zip(units, units[1:]):
    assert upper["base"] == lower["top"], f"{upper['code']} base and {lower['code']} top must be the same surface"
assert units[-1]["base"] is None, "the bottom unit runs down to the flat model base"
CREDIT = (meta.get("credit") or "BGS AGS data © UKRI; LIDAR © Environment Agency (both OGL v3).").rstrip()
CREDIT += f" Print vertical exaggeration x{VE:g}."

grids = []
for u in units:
    with rasterio.open(args.model / f"{u['top']}.tif") as src:
        win = from_bounds(*WINDOW, transform=src.transform).round_offsets().round_lengths()
        win = win.intersection(Window(0, 0, src.width, src.height))
        grids.append(src.read(1, window=win).astype("float64"))
        res, wt = src.res[0], src.window_transform(win)
S = np.stack(grids)                                   # (surface, row, col), top surface first
valid = np.isfinite(S).all(axis=0)
rows, cols = valid.shape
print(f"window {wt.c:.0f}-{wt.c + cols * res:.0f} E, {wt.f - rows * res:.0f}-{wt.f:.0f} N: {cols} x {rows} cells of "
      f"{res:.0f} m | model domain covers {valid.mean() * 100:.1f}% of it")
upside_down = int((np.diff(S, axis=0)[:, valid] > 1e-3).any(axis=0).sum())
S = np.minimum.accumulate(S, axis=0)                  # defensive: surfaces may only go down
below = int((S[:, valid] < base_m).any(axis=0).sum())
S = np.maximum(S, base_m)
print(f"{len(units)} units, model base {base_m:.0f} m OD | cells needing re-clamping {upside_down}, "
      f"below the model base (raised to it) {below}")
for u, s in zip(units, S):
    print(f"  top of {u['code']:<5s} {u['name']:<26s} {s[valid].min():6.1f} to {s[valid].max():6.1f} m OD")
# Cells outside the domain take their nearest valid value; they are cut away in Step 2, so nothing is invented.
near = ndimage.distance_transform_edt(~valid, return_distances=False, return_indices=True)
S = S[:, near[0], near[1]]

# ------------------------------------------------------------------ Step 2: footprint
section("Step 2: the footprint of the model domain")
scale = LONG_MM / ((max(rows, cols) - 1) * res)      # mm on the print per metre on the ground
cell = res * scale
# Mesh vertices sit on cell centres, so a grid square between four vertices is only real data if all four are
# valid. Polygonise those squares, shave one cell off the edge and simplify by no more than that: this turns the
# pixel staircase into straight walls without ever reaching outside the real data. Keep the largest piece.
quad = valid[:-1, :-1] & valid[1:, :-1] & valid[:-1, 1:] & valid[1:, 1:]
quad_to_mm = rasterio.Affine(cell, 0, 0, 0, -cell, (rows - 1) * cell)
foot = unary_union([shape(g) for g, _ in features.shapes(quad.astype("uint8"), mask=quad, transform=quad_to_mm)])
foot = foot.buffer(-cell, join_style="mitre").simplify(cell)
pieces = sorted(getattr(foot, "geoms", [foot]), key=lambda g: g.area, reverse=True)
foot = orient(pieces[0], 1.0)                         # outer ring anticlockwise, holes clockwise
print(f"horizontal scale 1:{1000 / scale:,.0f} | one {res:.0f} m cell = {cell:.2f} mm on the print")
print(f"footprint {foot.area / scale ** 2 / 1e6:.2f} km2 = {foot.area / 100:.1f} cm2 on the print, "
      f"{len(foot.interiors)} hole(s) | {len(pieces) - 1} smaller piece(s) left out "
      f"({sum(p.area for p in pieces[1:]) / scale ** 2:,.0f} m2)")
gx, gy = np.meshgrid(np.arange(cols) * cell, (rows - 1 - np.arange(rows)) * cell)
in_foot = valid & shapely.contains_xy(foot, gx, gy)
contours = [np.array(foot.exterior.coords)[:-1]] + [np.array(h.coords)[:-1] for h in foot.interiors]

# ------------------------------------------------------------------ Step 3: print heights and thin units
section(f"Step 3: print heights, vertical exaggeration x{VE:g}, nothing thinner than {MIN_MM} mm")
Z = (S - base_m) * scale * VE + BASE_MM
t_before = {}                                         # print thickness of each unit before the rule, for figure 03
print(f"1 m of ground = {scale * VE:.3f} mm on the print, so {MIN_MM} mm = {MIN_MM / (scale * VE):.2f} m of strata")
for k, u in enumerate(units[:-1]):
    t = Z[k] - Z[k + 1]
    thin = (t > 1e-6) & (t < MIN_MM)
    # Hand the sliver to the unit below: raise this unit's base (and any eroded, zero-thickness surfaces that
    # coincide with it) up to its top, so the next real unit down fills the space.
    old = Z[k + 1].copy()
    for j in range(k + 1, len(units)):
        Z[j] = np.where(thin & (Z[j] >= old - 1e-9), Z[k], Z[j])
    n_thin, present = int((thin & in_foot).sum()), int(((t > 1e-6) & in_foot).sum())
    print(f"{u['code']:<5s} present under {present / in_foot.sum() * 100:5.1f}% of the footprint | thinner than "
          f"{MIN_MM} mm under {n_thin / in_foot.sum() * 100:5.1f}% ({n_thin * res ** 2:,.0f} m2) -> absorbed by the unit below")
    t_before[u["code"]] = t[in_foot & (t > 1e-6)]

# ------------------------------------------------------------------ Step 4: one solid per unit
section("Step 4: cumulative solids and booleans")
prism = mf.Manifold.extrude(mf.CrossSection(contours), float(Z.max()) + 2).translate((0, 0, -1))
cum = [to_manifold(*heightfield_solid(z, cell)) for z in Z]      # plinth bottom up to each surface
# Subtract first, then trim to the footprint: trimming first would give every solid the same side walls, and
# subtracting coincident walls leaves needle-thin triangles on the outside of the print.
raw = [(cum[k] - cum[k + 1]) ^ prism for k in range(len(units) - 1)] + [cum[-1] ^ prism]
whole = cum[0] ^ prism
bx0, by0, _, bx1, by1, bz1 = whole.bounding_box()
print(f"whole model {bx1 - bx0:.1f} x {by1 - by0:.1f} x {bz1:.1f} mm, including the {BASE_MM:.0f} mm plinth")
solids = {}
for u, m in zip(units, raw):
    solids[u["code"]], n = drop_slivers(m)
    print(f"{u['code']:<5s} {n:>3d} separate piece(s)")

# ------------------------------------------------------------------ Step 5: checks
section("Step 5: check every part before writing")
meshes = {}
for u in units:
    m = solids[u["code"]]
    if m.is_empty():
        print(f"{u['code']:<5s} not present in the print window - no file")
        continue
    v, t = from_manifold(m)
    tight, oriented, vol = check_mesh(v, t)
    if not (tight and oriented and vol > 0):
        raise RuntimeError(f"{u['code']} is not printable: watertight {tight}, oriented {oriented}, volume {vol:.2f}")
    meshes[u["code"]] = (v, t, vol, tight)
codes = list(meshes)
total, whole = sum(m[2] for m in meshes.values()), whole.volume()
overlap = sum((solids[a] ^ solids[b]).volume() for i, a in enumerate(codes) for b in codes[i + 1:])
print(f"sum of parts {total / 1000:.3f} cm3 | ground solid {whole / 1000:.3f} cm3 | difference "
      f"{(total - whole) / whole * 100:+.4f}% | pairwise overlap {abs(overlap):.4f} mm3")
if abs(total - whole) > 0.001 * whole or overlap > 0.001 * whole:
    raise RuntimeError("the parts do not add up to the ground solid")

# ------------------------------------------------------------------ Step 6: files
section("Step 6: write one STL per unit and a colour 3MF")
UNIT = {u["code"]: u for u in units}
stl = {}
print(f"{'file':<14s} {'unit':<26s} {'triangles':>9s} {'cm3':>7s} {'PLA g':>6s}  watertight")
for i, code in enumerate(codes):
    v, t, vol, tight = meshes[code]
    stl[code] = OUT / f"{i + 1:02d}_{code}.stl"
    write_stl(stl[code], v, t, f"TechnicalGIS ground model {code} - BGS AGS, EA LIDAR (OGL v3)")
    print(f"{stl[code].name:<14s} {UNIT[code]['name']:<26s} {len(t):>9,} {vol / 1000:7.2f} "
          f"{vol / 1000 * PLA_G_CM3:6.1f}  {tight}")
print(f"{'total':<41s} {sum(len(m[1]) for m in meshes.values()):>9,} {total / 1000:7.2f} {total / 1000 * PLA_G_CM3:6.1f}"
      f"  (solid PLA; at 15% infill a slicer will use much less)")
write_3mf(OUT / "ground_model.3mf", [(UNIT[c]["name"], *meshes[c][:2], UNIT[c]["colour"]) for c in codes])
with zipfile.ZipFile(OUT / "ground_model.3mf") as z:
    objs = ET.fromstring(z.read("3D/3dmodel.model")).findall(
        ".//{http://schemas.microsoft.com/3dmanufacturing/core/2015/02}object")
print(f"ground_model.3mf {(OUT / 'ground_model.3mf').stat().st_size / 1e6:.1f} MB | {len(objs)} parts: "
      + ", ".join(o.get("name") for o in objs))

# ------------------------------------------------------------------ Step 7: previews from the STLs on disk
section("Step 7: previews rendered from the STL files")
parts = {c: pv.read(stl[c]) for c in codes}
x0, x1, y0, y1, _, z1 = pv.merge(list(parts.values())).bounds
print(f"read back {len(parts)} STLs: {x1 - x0:.1f} x {y1 - y0:.1f} x {z1:.1f} mm")


def compose(img, path, title, note=CREDIT):
    """Put a PyVista screenshot beside a proper legend (PyVista's own legend tints the text)."""
    ink = np.where((img[..., :3] < 245).any(axis=2))
    img = img[max(ink[0].min() - 10, 0):ink[0].max() + 10, max(ink[1].min() - 10, 0):ink[1].max() + 10]
    h, w = img.shape[:2]
    fig = plt.figure(figsize=(11, 11 * h / w * 0.86))
    ax = fig.add_axes([0, 0.035, 0.8, 0.905])
    ax.imshow(img)
    ax.set_axis_off()
    fig.text(0.01, 0.975, title, fontsize=11, weight="bold", va="top")
    fig.text(0.01, 0.005, note, fontsize=7, color="#444444")
    handles = [Patch(facecolor=UNIT[c]["colour"], edgecolor="#555555", lw=0.4, label=f"{i + 1:02d} {UNIT[c]['name']}")
               for i, c in enumerate(codes)]
    fig.legend(handles=handles, loc="center left", bbox_to_anchor=(0.79, 0.5), fontsize=7.5, frameon=False,
               title="Part (STL)", title_fontsize=8, alignment="left")
    fig.savefig(OUT / path, dpi=150)
    plt.close(fig)


def render(offsets=None, elev=28, azim=-125, dist=1.9, top=False):
    p = pv.Plotter(off_screen=True, window_size=(1800, 1200))
    p.set_background("white")
    for i, c in enumerate(codes):
        mesh = parts[c].translate((0, 0, offsets[i]), inplace=False) if offsets else parts[c]
        p.add_mesh(mesh, color=UNIT[c]["colour"], ambient=0.3, diffuse=0.7, specular=0.05)
    p.enable_lightkit()
    centre = np.array(((x0 + x1) / 2, (y0 + y1) / 2, (offsets[0] if offsets else 0) / 2 + z1 / 3))
    if top:
        p.view_xy()
        p.enable_parallel_projection()
    else:
        el, az = np.radians(elev), np.radians(azim)
        eye = centre + dist * max(x1 - x0, y1 - y0) * np.array([np.cos(el) * np.cos(az), np.cos(el) * np.sin(az), np.sin(el)])
        p.camera_position = [tuple(eye), tuple(centre), (0, 0, 1)]
        p.camera.view_angle = 30
    img = p.screenshot(return_img=True)
    p.close()
    return img


size = f"{x1 - x0:.0f} x {y1 - y0:.0f} x {z1:.0f} mm"
compose(render(), "01-printable-ground-model.png",
        f"Crickley Hill ground model as a {len(codes)}-part print ({size}), looking north-east up the escarpment")
compose(render(offsets=[(len(codes) - 1 - i) * EXPLODE_MM for i in range(len(codes))], dist=2.1),
        "02-exploded-view.png", f"The same {len(codes)} STLs pulled apart by {EXPLODE_MM:.0f} mm: each unit is its own solid")
compose(render(top=True), "04-slicer-style-top.png",
        "Seen from above (north up), the print is a geological map: where each unit reaches the surface")

# Section through the STLs: turn each part so that a horizontal Manifold slice cuts it along a west-east line.
row = int(np.argmax(quad.sum(axis=1)))
y_cut = (rows - 1.5 - row) * cell
fig, (ax, bx) = plt.subplots(2, 1, figsize=(10, 6.4), gridspec_kw={"height_ratios": [1.3, 1]})
for c in codes:
    faces = parts[c].faces.reshape(-1, 4)[:, 1:]
    cs = to_manifold(parts[c].points, faces).rotate((-90, 0, 0)).slice(-y_cut)
    polys = [np.vstack([p, p[:1]]) for p in cs.to_polygons()]
    if polys:
        verts = np.vstack(polys)
        codes_ = np.concatenate([[MplPath.MOVETO] + [MplPath.LINETO] * (len(p) - 2) + [MplPath.CLOSEPOLY] for p in polys])
        ax.add_patch(PathPatch(MplPath(verts, codes_), facecolor=UNIT[c]["colour"], edgecolor="#333333", lw=0.4))
ax.autoscale_view()
ax.set_aspect("equal")
ax.set_xlabel("mm along the print (west to east)")
ax.set_ylabel("mm above the bed")
ax.set_title(f"a) Section through the STLs {y_cut:.1f} mm from the front edge "
             f"({wt.f - (row + 1) * res:.0f} m N), true print scale", fontsize=9, loc="left")
for i, c in enumerate(codes[:-1]):
    th = t_before.get(c, np.array([]))
    if th.size:
        bx.hist(th, bins=np.arange(0, np.percentile(th, 99) + 0.2, 0.2), histtype="step", lw=1.4,
                color=UNIT[c]["colour"], label=UNIT[c]["name"])
bx.axvline(MIN_MM, color="#b2182b", lw=1.2, ls="--")
bx.text(MIN_MM, bx.get_ylim()[1] * 0.95, f"  {MIN_MM} mm: thinner is absorbed by the unit below", color="#b2182b",
        fontsize=7.5, va="top")
bx.set_xlabel("unit thickness on the print (mm)")
bx.set_ylabel(f"{res:.0f} m cells")
bx.set_title("b) Where each unit is present, how thick it prints", fontsize=9, loc="left")
bx.legend(fontsize=7, frameon=False)
for a in (ax, bx):
    a.spines[["top", "right"]].set_visible(False)
fig.text(0.01, -0.01, CREDIT, fontsize=7, color="#444444")
fig.tight_layout()
fig.savefig(OUT / "03-print-layers-section.png")
plt.close(fig)
print("figures written:", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))
print(f"finished in {time.perf_counter() - T0:.1f} s")