TechnicalGIS

Source code: A 3D ground model from AGS4 boreholes: kriged stratum surfaces in Python

Runnable Python scripts, requirements and console output for the tutorial “A 3D ground model from AGS4 boreholes: kriged stratum surfaces in Python”.

Last updated

These are the exact scripts used to produce every figure and number in A 3D ground model from AGS4 boreholes: kriged stratum surfaces in Python. Download them, install the requirements into a fresh virtual environment and run them from this folder.

Run it
python -m venv .venv
.venv\Scripts\activate      # macOS/Linux: source .venv/bin/activate
pip install -r requirements.txt
python a417_ags.py

a417_ags.py

a417_ags.py
"""Download, merge and repair the A417 Missing Link ground investigation AGS4 data.


Shared by every part of the TechnicalGIS series "AGS data in GIS"; each part keeps an identical copy.


The British Geological Survey (BGS) holds AGS files donated by the ground investigation industry.
Its OGC API lists every borehole it can export (collection "agsboreholeindex"), and its AGS
File Utilities API exports those boreholes as AGS4 files, one ZIP per request. A request URL
can only carry so many borehole IDs, so we ask for them in batches of 40.


repair() applies the fixes worked through step by step in Part 1, so Parts 2 and 3 can start
from a clean, merged set of tables with load_clean().


Data: A417 Missing Link ground investigation (BGS deposited data item 173984; client Highways
England, contractor Geotechnical Engineering Ltd), from the BGS AGS download service under the
Open Government Licence v3.0. Contains British Geological Survey materials © UKRI 2026.
"""
import io
import time
import zipfile
from pathlib import Path


import pandas as pd
import requests
from python_ags4 import AGS4


ITEM = 173984                       # BGS deposited data item: HE551505 A417 Missing Link GI
INDEX = "https://ogcapi.bgs.ac.uk/collections/agsboreholeindex/items"
EXPORT = "https://agsapi.bgs.ac.uk/v7/ags_export/"
BATCH = 40
UA = {"User-Agent": "TechnicalGIS-tutorial (https://www.technicalgis.com)"}


# Geological units met in the A417 logs, youngest at the top. Bedrock codes are BGS Lexicon codes
# (GEOL_BGS); rank gives the expected order down a borehole, and an undivided group spans the ranks of its
# formations (rank to last). Superficial deposits have rank 0.
UNITS = pd.DataFrame([
    # code,  name,                                   group,                   rank, colour
    ("TOPS", "Topsoil",                              "Superficial",            0, "#5b4636"),
    ("MGR",  "Made ground",                          "Superficial",            0, "#9e9e9e"),
    ("HEAD", "Head",                                 "Superficial",            0, "#d9b77e"),
    ("SLIP", "Landslide deposits",                   "Superficial",            0, "#c98c5a"),
    ("WHL",  "White Limestone Formation",            "Great Oolite Group",     1, "#f7f1bd"),
    ("HMB",  "Hampen Formation",                     "Great Oolite Group",     2, "#e4d79c"),
    ("GOG",  "Great Oolite Group (undivided)",       "Great Oolite Group",     1, "#eee3a6"),
    ("FE",   "Fuller's Earth Formation",             "Great Oolite Group",     3, "#9fb08f"),
    ("SALS", "Salperton Limestone Formation",        "Inferior Oolite Group",  4, "#f4cd72"),
    ("ASLS", "Aston Limestone Formation",            "Inferior Oolite Group",  5, "#e9b157"),
    ("BLPL", "Birdlip Limestone Formation",          "Inferior Oolite Group",  6, "#d98f3e"),
    ("IOG",  "Inferior Oolite Group (undivided)",    "Inferior Oolite Group",  4, "#e6a855"),
    ("BDS",  "Bridport Sand Formation",              "Lias Group",             7, "#d2c48c"),
    ("WHM",  "Whitby Mudstone Formation",            "Lias Group",             8, "#7d93ab"),
    ("MRB",  "Marlstone Rock Formation",             "Lias Group",             9, "#b0703f"),
    ("DYS",  "Dyrham Formation",                     "Lias Group",            10, "#a3b8cc"),
    ("CHAM", "Charmouth Mudstone Formation",         "Lias Group",            11, "#55698a"),
    ("NR",   "No recovery or not classified",        "Other",                 0, "#f2f2f2"),
], columns=["code", "name", "group", "rank", "colour"]).set_index("code")
UNITS["last"] = UNITS["rank"]
UNITS.loc[["GOG", "IOG"], "last"] = [3, 6]


# Where GEOL_BGS is blank, the project's own code (GEOL_GEOL) still tells us the unit.
FROM_GEOL = {"TOP": "TOPS", "MG": "MGR", "HDD": "HEAD", "GOG": "GOG", "IOG": "IOG", "SLIP": "SLIP"}




def borehole_index(item=ITEM):
    """Every borehole in one BGS deposited data item, from the BGS OGC API (GeoJSON)."""
    r = requests.get(INDEX, params={"f": "json", "limit": 1000, "item_index_id": item}, headers=UA, timeout=120)
    r.raise_for_status()
    rows = [f["properties"] for f in r.json()["features"]]
    return pd.DataFrame(rows).sort_values("bgs_loca_id").reset_index(drop=True)




def _get(url, params, tries=5):
    for attempt in range(tries):
        try:
            r = requests.get(url, params=params, headers=UA, timeout=300)
            r.raise_for_status()
            return r
        except requests.RequestException:
            if attempt == tries - 1:
                raise
            time.sleep(10 * (attempt + 1))




def download(data_dir, item=ITEM):
    """Download the item's boreholes as AGS4 files, once. Returns the .ags paths in order."""
    raw = Path(data_dir) / f"bgs_{item}"
    raw.mkdir(parents=True, exist_ok=True)
    if not list(raw.glob("*.ags")):
        ids = borehole_index(item).bgs_loca_id.tolist()
        for i in range(0, len(ids), BATCH):
            r = _get(EXPORT, {"bgs_loca_id": ";".join(ids[i:i + BATCH])})
            with zipfile.ZipFile(io.BytesIO(r.content)) as z:
                for name in z.namelist():
                    (raw / f"batch{i // BATCH + 1:02d}_{Path(name).name}").write_bytes(z.read(name))
            time.sleep(1)
    return sorted(raw.glob("*.ags"))




def merge(paths):
    """Read several AGS4 files from one project and stack them, group by group.


    Each table keeps its UNIT and TYPE rows (from the first file) followed by every DATA row;
    rows repeated in every file (PROJ, TRAN, ABBR, DICT and so on) are kept once.
    """
    tables, headings = {}, {}
    for p in paths:
        t, h = AGS4.AGS4_to_dataframe(str(p))
        for group, df in t.items():
            if group not in tables:
                tables[group], headings[group] = df, list(h[group])
                continue
            headings[group] += [c for c in h[group] if c not in headings[group]]
            tables[group] = pd.concat([tables[group], df[df.HEADING == "DATA"]], ignore_index=True)
    for group, df in tables.items():
        df = df.reindex(columns=headings[group]).fillna("")
        tables[group] = df.drop_duplicates(ignore_index=True)
    return tables, headings




def repair(tables, headings):
    """Apply the Part 1 fixes to merged A417 tables, in place, and return them."""
    # Rule 1: a degree sign double-encoded as "Β°" (Greek capital beta + degree); the logs write "deg".
    for df in tables.values():
        df.replace("Β°", "deg", regex=True, inplace=True)


    # Rule 10a: three DICT headings defined twice; keep the later, archive descriptions.
    d = tables["DICT"]
    tables["DICT"] = d[~d.duplicated(["HEADING", "DICT_TYPE", "DICT_GRP", "DICT_HDNG"], keep="last")
                       | (d.HEADING != "DATA")].reset_index(drop=True)


    # Rules 7, 9, 19a, 19b: rename the BGS identifiers to valid user-defined headings and define them.
    new_dict = []
    for group, old, new, desc in [("PROJ", "BGS_PROJ_ID", "PROJ_BGSI", "BGS export project identifier"),
                                  ("LOCA", "BGS_LOCA_ID", "LOCA_BGSI", "BGS borehole index identifier")]:
        tables[group] = tables[group].rename(columns={old: new})
        tables[group].loc[tables[group].HEADING == "TYPE", new] = "X"
        headings[group] = [new if c == old else c for c in headings[group]]
        new_dict.append({"HEADING": "DATA", "DICT_TYPE": "HEADING", "DICT_GRP": group, "DICT_HDNG": new,
                         "DICT_STAT": "OTHER", "DICT_DTYP": "X", "DICT_DESC": desc})


    # The *_TEXT columns hold the depths and diameters exactly as originally written ("0.50"), where the
    # plain columns were reformatted ("0.5"). Put the original text back, restore the standard data type
    # and drop the extra column.
    std, _ = AGS4.AGS4_to_dataframe(str(Path(AGS4.__file__).parent / "Standard_dictionary_v4_0_3.ags"))
    dtype = std["DICT"].query("DICT_TYPE == 'HEADING'").set_index(["DICT_GRP", "DICT_HDNG"]).DICT_DTYP
    for group in list(tables):
        df = tables[group]
        for col in [c for c in headings[group] if c.endswith("_TEXT")]:
            base = col[:-5]
            data = df.HEADING == "DATA"
            same = pd.to_numeric(df.loc[data, base]).round(6).equals(pd.to_numeric(df.loc[data, col]).round(6))
            assert same, f"{group}.{col} disagrees with {base}"
            df.loc[data, base] = df.loc[data, col]
            df.loc[df.HEADING == "TYPE", base] = dtype[(group, base)]
            tables[group] = df = df.drop(columns=col)
            headings[group].remove(col)


    # Rule 10c: in the child groups LOCA_ID was exported with TYPE "X", so it no longer matches the
    # parent key in LOCA (TYPE "ID"). Restore the standard type.
    for group, df in tables.items():
        if group != "LOCA" and "LOCA_ID" in df:
            df.loc[df.HEADING == "TYPE", "LOCA_ID"] = "ID"


    tables["DICT"] = pd.concat([tables["DICT"], pd.DataFrame(new_dict)], ignore_index=True).fillna("")


    # Rule 7 (only checkable once the headings above are defined): headings must follow the order of
    # the standard dictionary, with the file's own DICT headings after them in DICT order.
    own = tables["DICT"].query("HEADING == 'DATA' and DICT_TYPE == 'HEADING'")
    defined = pd.concat([std["DICT"].query("DICT_TYPE == 'HEADING'"), own])
    for group in tables:
        listed = defined[defined.DICT_GRP == group].DICT_HDNG.drop_duplicates().tolist()
        order = {h: i for i, h in enumerate(listed)}
        headings[group] = ["HEADING"] + sorted(headings[group][1:], key=lambda h: order.get(h, len(order)))
        tables[group] = tables[group][headings[group]]


    tables["TRAN"].loc[tables["TRAN"].HEADING == "DATA", "TRAN_DESC"] += \
        "; three BGS exports merged and repaired (TechnicalGIS AGS data in GIS, Part 1)"
    return tables, headings




def write(tables, headings, ags_path, metadata):
    """Write a clean AGS4 file with the FILE folder AGS Rule 20 expects beside it."""
    ags_path = Path(ags_path)
    ags_path.parent.mkdir(parents=True, exist_ok=True)
    AGS4.dataframe_to_AGS4(tables, headings, str(ags_path), warnings=False)
    fset = ags_path.parent / "FILE" / "BGSFileSet01"
    fset.mkdir(parents=True, exist_ok=True)
    (fset / "BGS_download_metadata.txt").write_bytes(Path(metadata).read_bytes())
    return ags_path




def numeric(tables):
    """DATA rows only, with numbers as numbers. python-ags4 reads the TYPE row to decide; it leaves
    "U" (variable decimal places) as text, so convert those columns too."""
    out = {}
    for g, df in tables.items():
        types = df[df.HEADING == "TYPE"].iloc[0]
        n = AGS4.convert_to_numeric(df)
        for col in types.index[types == "U"]:
            n[col] = pd.to_numeric(n[col], errors="coerce")
        out[g] = n
    return out




def unit_code(geol):
    """One unit code per GEOL row: the BGS Lexicon code, else the project code, else NR."""
    code = geol.GEOL_BGS.where(geol.GEOL_BGS != "", geol.GEOL_GEOL.map(FROM_GEOL))
    return code.fillna("NR")




def in_order(codes):
    """True if a top-to-bottom list of bedrock unit codes never steps back to a younger unit."""
    reached = 0
    for code in codes:
        if UNITS.at[code, "last"] < reached:
            return False
        reached = max(reached, UNITS.at[code, "rank"])
    return True




def load_clean(data_dir):
    """Download (once), merge and repair the A417 data. Returns numeric DataFrames by group."""
    tables, headings = repair(*merge(download(data_dir)))
    out = numeric(tables)
    out["GEOL"] = out["GEOL"].sort_values(["LOCA_ID", "GEOL_TOP"], ignore_index=True)
    out["GEOL"]["UNIT"] = unit_code(out["GEOL"])
    return out

ground_model.py

ground_model.py
"""Interpolated stratum surfaces from AGS4: a 3D ground model of the A417 escarpment, with its evidence.


Part 5 of the TechnicalGIS series "AGS data in GIS". Picks four geological surfaces from the A417
Missing Link borehole logs (rockhead, the top and base of the Inferior Oolite and the base of the
Bridport Sand), kriges them onto a 10 m grid under Environment Agency LIDAR ground, and then tests the
result: leave-one-out cross-validation at the boreholes, kriging standard deviation maps, every borehole
that stops short of a surface used as a one-sided check, and a check that no surface crosses another.
The two places Part 3 found where the logs disagree with their neighbours (the slipped block in CP217 and
the low block at RC507-RC509) are handled explicitly rather than smoothed away.


TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/geological-ground-model-surfaces-python/


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 json
import time
from importlib.metadata import version
from pathlib import Path


import numpy as np
import pandas as pd
import requests
import rasterio
from rasterio.transform import from_origin
from scipy.optimize import curve_fit
from scipy.spatial import cKDTree
from pykrige.ok import OrdinaryKriging
import pyvista as pv
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LightSource
from matplotlib.lines import Line2D
from matplotlib.patches import Patch


import a417_ags as A


plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9,
                     "figure.facecolor": "white", "axes.facecolor": "white"})
pd.set_option("display.width", 130)


DATA = Path("data"); DATA.mkdir(exist_ok=True)
OUT = Path("output"); OUT.mkdir(exist_ok=True)
MODEL = OUT / "model"; MODEL.mkdir(exist_ok=True)
UA = {"User-Agent": "TechnicalGIS-tutorial (https://www.technicalgis.com)"}


ITEMS = [173984, 173975]                                # the main GI, plus 20 extra holes on the escarpment crest
AOI = (391_600, 212_800, 395_600, 216_800)             # same 4 x 4 km window as Parts 2 and 3
RES = 5.0                                               # LIDAR download resolution (m)
CELL = 10.0                                             # model grid (m)
REACH = 300.0                                           # model only within this distance of a borehole
VE = 3.0
SLIPPED = {"CP217": 19.10}                              # Part 3: a slipped block repeats the Lias above this depth
CAMBER = ["RC507", "RC508", "RC509"]                    # Part 3: base of Inferior Oolite 11-15 m below neighbours
WCS = ("https://environment.data.gov.uk/spatialdata/lidar-composite-digital-terrain-model-dtm-1m/wcs"
       "?service=WCS&version=2.0.1&request=GetCoverage&format=image/tiff"
       "&CoverageId=13787b9a-26a4-4775-8523-806d13af58fc__Lidar_Composite_Elevation_DTM_1m")
CREDIT = "BGS AGS data © UKRI; LIDAR © Environment Agency (both OGL v3)."


# Model units, youngest first, and the GEOL units that belong to each.
GROUPS = {"SUP": ["TOPS", "MGR", "HEAD", "SLIP"], "GOG": ["WHL", "HMB", "GOG", "FE"],
          "IOG": ["SALS", "ASLS", "BLPL", "IOG"], "BDS": ["BDS"], "LIAS": ["WHM", "MRB", "DYS", "CHAM"]}
ORDER = list(GROUPS)
GROUP = {u: g for g, us in GROUPS.items() for u in us}
UNIT = pd.DataFrame({
    "name": ["Superficial deposits", "Great Oolite Group", "Inferior Oolite Group", "Bridport Sand Formation",
             "Lias mudstones (Whitby Mudstone and older)"],
    "colour": ["#c98c5a", "#f7f1bd", "#d98f3e", "#bfae6e", "#7d93ab"]}, index=ORDER)
# Each modelled surface is the base of one group and the top of the next.
SURFACES = {"rockhead": "base of the superficial deposits", "top_io": "top of the Inferior Oolite",
            "base_io": "base of the Inferior Oolite", "base_bds": "base of the Bridport Sand"}
BELOW = {"top_io": ("GOG", "IOG"), "base_io": ("IOG", "BDS"), "base_bds": ("BDS", "LIAS")}




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




def get(url, path, tries=6):
    """Download once and cache, backing off when the Defra WCS is busy."""
    if not path.exists():
        for attempt in range(tries):
            try:
                r = requests.get(url, timeout=300, headers=UA)
                r.raise_for_status()
                if not r.headers.get("Content-Type", "").startswith("image/"):
                    raise requests.RequestException(f"not an image: {r.text[:200]}")
                break
            except requests.RequestException as err:
                if attempt == tries - 1:
                    raise
                wait = getattr(err.response, "headers", {}).get("Retry-After")
                time.sleep(float(wait) if wait and wait.isdigit() else 20 * (attempt + 1))
        path.write_bytes(r.content)
    return path




print(f"pykrige {version('pykrige')} | pyvista {pv.__version__} | numpy {np.__version__} | "
      f"python-ags4 {version('python-ags4')}")


# ------------------------------------------------------------------ Step 1: logs to model units
section("Step 1: logs to model units")
paths = [p for item in ITEMS for p in A.download(DATA, item)]
tables, headings = A.repair(*A.merge(paths))
ags = A.numeric(tables)
loca = ags["LOCA"].set_index("LOCA_ID")
geol = ags["GEOL"].sort_values(["LOCA_ID", "GEOL_TOP"], ignore_index=True)
geol["UNIT"] = A.unit_code(geol)
geol = geol[geol.GEOL_BASE > geol.GEOL_TOP].join(loca[["LOCA_NATE", "LOCA_NATN", "LOCA_GL"]], on="LOCA_ID")
geol["group"] = geol.UNIT.map(GROUP)                     # NR (no recovery) stays blank and is skipped
for lid, depth in SLIPPED.items():
    moved = (geol.LOCA_ID == lid) & (geol.GEOL_TOP < depth) & geol.group.notna()
    print(f"{lid}: {moved.sum()} intervals above {depth} m (a repeated Lias sequence) counted with the landslide mass")
    geol.loc[moved, "group"] = "SUP"
print(f"{loca.shape[0]} holes from items {', '.join(map(str, ITEMS))}; "
      f"{geol.group.notna().groupby(geol.LOCA_ID).any().sum()} have at least one classified interval")


# Collapse each log to a run of model units, top to bottom.
runs = []
for lid, g in geol[geol.group.notna()].groupby("LOCA_ID", sort=False):
    g = g.assign(run=(g.group != g.group.shift()).cumsum())
    r = g.groupby("run").agg(group=("group", "first"), top=("GEOL_TOP", "min"), base=("GEOL_BASE", "max"))
    runs.append(r.assign(LOCA_ID=lid))
runs = pd.concat(runs, ignore_index=True).join(loca[["LOCA_NATE", "LOCA_NATN", "LOCA_GL", "LOCA_FDEP"]],
                                               on="LOCA_ID")
runs["rank"] = runs.group.map(ORDER.index)
backwards = runs.groupby("LOCA_ID")["rank"].apply(lambda s: (s.diff() < 0).any())
print("holes whose model units step back to a younger unit:", ", ".join(backwards[backwards].index) or "none")


# ------------------------------------------------------------------ Step 2: pick the contacts
section("Step 2: pick the contacts, and what the other holes say")
picks, bounds = [], []
for lid, r in runs.groupby("LOCA_ID", sort=False):
    h = r.iloc[0]
    xy = dict(LOCA_ID=lid, E=h.LOCA_NATE, N=h.LOCA_NATN, GL=h.LOCA_GL)
    ranks, bottom = r["rank"].to_numpy(), h.LOCA_GL - r.base.max()
    # Rockhead: the first bedrock run. A hole that never leaves the superficial deposits only tells us
    # the superficial deposits are at least that thick.
    rock = r[r.group != "SUP"]
    if rock.empty:
        bounds.append(dict(xy, surface="rockhead", kind="deeper than", z=bottom))
    else:
        sup = r[r.group == "SUP"]
        depth = (sup.base.max() + rock.top.iloc[0]) / 2 if len(sup) else rock.top.iloc[0]
        picks.append(dict(xy, surface="rockhead", depth=depth, z=h.LOCA_GL - depth))
    first_rock = rock["rank"].iloc[0] if len(rock) else None
    for s, (up, down) in BELOW.items():
        a, b = r[r.group == up], r[r.group == down]
        if len(a) and len(b):                           # seen in the hole: the midpoint of any unlogged gap
            depth = (a.base.max() + b.top.min()) / 2
            picks.append(dict(xy, surface=s, depth=depth, z=h.LOCA_GL - depth))
        elif first_rock is not None and first_rock >= ORDER.index(down):
            # the hole starts in older rock, so this surface has been eroded and lay above rockhead
            bounds.append(dict(xy, surface=s, kind="above", z=h.LOCA_GL - rock.top.iloc[0]))
        elif len(rock) and ranks.max() <= ORDER.index(up):
            bounds.append(dict(xy, surface=s, kind="below", z=bottom))   # stops in older bedrock above it
picks, bounds = pd.DataFrame(picks), pd.DataFrame(bounds)
picks["camber"] = picks.LOCA_ID.isin(CAMBER)
print("contacts seen in the logs (holes):", picks.groupby("surface").size().reindex(SURFACES).to_dict())
print("one-sided checks from holes that miss a surface:")
print(bounds.groupby(["surface", "kind"]).size().unstack(fill_value=0).reindex(SURFACES).to_string())
print(picks[picks.surface != "rockhead"].round(2).sort_values(["surface", "E"]).to_string(index=False))


# ------------------------------------------------------------------ Step 3: ground and the model grid
section("Step 3: LIDAR ground and the model grid")
X0, Y0, X1, Y1 = AOI
f = get(f"{WCS}&subset=E({X0},{X1})&subset=N({Y0},{Y1})&SCALEFACTOR={1 / RES}",
        DATA / f"ea_dtm_{X0}_{Y0}_{X1 - X0}m_{RES:.0f}m.tif")
with rasterio.open(f) as src:
    dtm5 = src.read(1).astype("float64")
    dtm5[dtm5 < -1000] = np.nan
k = int(CELL / RES)
ground = dtm5.reshape(dtm5.shape[0] // k, k, dtm5.shape[1] // k, k).mean(axis=(1, 3))   # 5 m to 10 m
ny, nx = ground.shape
transform = from_origin(X0, Y1, CELL, CELL)
gx, gy = np.meshgrid(CELL / 2 + np.arange(nx) * CELL, (Y1 - Y0) - CELL / 2 - np.arange(ny) * CELL)  # local metres
holes = runs.drop_duplicates("LOCA_ID")
hx, hy = holes.LOCA_NATE.to_numpy() - X0, holes.LOCA_NATN.to_numpy() - Y0
dist, _ = cKDTree(np.c_[hx, hy]).query(np.c_[gx.ravel(), gy.ravel()])
inside = (dist.reshape(ny, nx) <= REACH) & np.isfinite(ground)
print(f"grid {nx} x {ny} cells of {CELL:.0f} m; {inside.sum():,} cells ({inside.sum() * CELL**2 / 1e6:.2f} km²) "
      f"lie within {REACH:.0f} m of one of the {len(holes)} logged holes")
picks["x"], picks["y"] = picks.E - X0, picks.N - Y0
bounds["x"], bounds["y"] = bounds.E - X0, bounds.N - Y0
row = np.clip(((Y1 - picks.N) // CELL).astype(int), 0, ny - 1)
col = np.clip(((picks.E - X0) // CELL).astype(int), 0, nx - 1)
gl_minus_dtm = (picks.GL - ground[row, col])[picks.surface == "rockhead"]
print(f"logged ground level minus LIDAR at the holes: median {gl_minus_dtm.median():+.2f} m, "
      f"90% within {gl_minus_dtm.abs().quantile(0.9):.2f} m")




def dedupe(p, tol=2.0):
    """Average holes drilled within tol metres of each other (a repeat hole would make kriging singular)."""
    p = p.reset_index(drop=True)
    label = np.arange(len(p))
    for i, j in sorted(cKDTree(p[["x", "y"]]).query_pairs(tol)):
        label[label == label[j]] = label[i]
    num = p.select_dtypes("number").columns
    out = p.groupby(label).agg({**{c: "mean" for c in num}, "LOCA_ID": "+".join, "surface": "first"})
    return out.reset_index(drop=True)




def experimental(x, y, v, edges):
    """Semivariance in distance bins: half the mean squared difference of every pair of holes."""
    i, j = np.triu_indices(len(v), 1)
    d = np.hypot(x[i] - x[j], y[i] - y[j])
    g = 0.5 * (v[i] - v[j]) ** 2
    b = np.digitize(d, edges) - 1
    ok = (b >= 0) & (b < len(edges) - 1)
    n = np.bincount(b[ok], minlength=len(edges) - 1)
    lag = np.bincount(b[ok], d[ok], len(edges) - 1) / np.maximum(n, 1)
    gam = np.bincount(b[ok], g[ok], len(edges) - 1) / np.maximum(n, 1)
    return lag[n > 0], gam[n > 0], n[n > 0]




def spherical(h, psill, rng, nugget):
    h = np.minimum(np.asarray(h, float) / rng, 1.0)
    return nugget + psill * (1.5 * h - 0.5 * h**3)




def fit_spherical(x, y, v, edges):
    """Fit a spherical variogram to the binned semivariances, weighting each bin by its number of pairs."""
    lag, gam, n = experimental(x, y, v, edges)
    p, _ = curve_fit(spherical, lag, gam, p0=[gam.max(), edges[-1] / 3, 0], sigma=1 / np.sqrt(n),
                     bounds=([0, edges[1], 0], [10 * gam.max(), edges[-1], gam.max()]))
    return dict(psill=p[0], range=p[1], nugget=p[2]), (lag, gam, n)




def krige(x, y, v, params, trend=False):
    """Ordinary kriging with a fixed spherical variogram, optionally of the residuals from a fitted plane."""
    X = np.c_[np.ones_like(x), x, y]
    coef = np.linalg.lstsq(X, v, rcond=None)[0] if trend else np.zeros(3)
    ok = OrdinaryKriging(x, y, v - X @ coef, variogram_model="spherical", variogram_parameters=params)


    def predict(px, py):
        r, ss = ok.execute("points", np.asarray(px, float), np.asarray(py, float))
        return np.asarray(r) + np.c_[np.ones_like(px), px, py] @ coef, np.sqrt(np.clip(np.asarray(ss), 0, None))
    return predict, coef




def residuals(x, y, v):
    X = np.c_[np.ones_like(x), x, y]
    coef = np.linalg.lstsq(X, v, rcond=None)[0]
    return v - X @ coef, coef




def loocv(x, y, v, params, trend=False, extra=((), (), ())):
    """Leave each hole out in turn, rebuild from the rest (same variogram), and predict it."""
    pred, sd = np.empty(len(v)), np.empty(len(v))
    for i in range(len(v)):
        keep = np.arange(len(v)) != i
        f, _ = krige(np.r_[x[keep], extra[0]], np.r_[y[keep], extra[1]], np.r_[v[keep], extra[2]], params, trend)
        p, s = f(x[i:i + 1], y[i:i + 1])
        pred[i], sd[i] = p[0], s[0]
    return pred, sd




TOL = 1.0                                               # metres of slack for logged levels and 10 m cells




def honour(name, x, y, v, params, b, trend=False):
    """Krige, then turn every one-sided check the surface breaks into a data point at its limit, and repeat.


    b has columns x, y, kind ('above' or 'below': the surface must lie above or below b.limit) and limit.
    """
    added = np.zeros(len(b), bool)
    first = None
    for _ in range(20):
        ex = b[added]
        f, coef = krige(np.r_[x, ex.x], np.r_[y, ex.y], np.r_[v, ex.limit], params, trend)
        m, _ = f(b.x.values, b.y.values)
        broken = np.where(b.kind == "above", m < b.limit - TOL, m > b.limit + TOL) & ~added
        first = broken.sum() if first is None else first
        if not broken.any():
            break
        added |= broken
    ex = b[added]
    passes[name] = (len(b), first, added.sum())
    print(f"  one-sided checks: the first pass broke {first} of {len(b)}; {added.sum()} holes added at their "
          f"limits; now broken: {broken.sum()}")
    return f, coef, (ex.x.values, ex.y.values, ex.limit.values), b.LOCA_ID[added].tolist()




def scores(obs, pred, sd):
    e = pred - obs
    return dict(n=len(e), RMSE=np.sqrt(np.mean(e**2)), mean_error=e.mean(), MAE=np.abs(e).mean(),
                mean_SE=np.mean(e / sd), var_SE=np.var(e / sd))




def on_grid(f):
    p, s = f(gx[inside], gy[inside])
    a, b = np.full(ground.shape, np.nan), np.full(ground.shape, np.nan)
    a[inside], b[inside] = p, s
    return a, b




cv, vgm, passes = {}, {}, {}
LAGS = np.arange(0, 2050, 100.0)                         # 100 m bins out to 2 km


# ------------------------------------------------------------------ Step 4: rockhead
section("Step 4: rockhead = LIDAR ground minus kriged superficial thickness")




def at(grid, E, N):
    """Value of a model grid at British National Grid points (nearest cell)."""
    r = np.clip(((Y1 - np.asarray(N)) // CELL).astype(int), 0, ny - 1)
    c = np.clip(((np.asarray(E) - X0) // CELL).astype(int), 0, nx - 1)
    return grid[r, c]




def show(name):
    p = vgm[name][0]
    print(f"  variogram: spherical, partial sill {p['psill']:.1f} m², range {p['range']:.0f} m, "
          f"nugget {p['nugget']:.1f} m²")




def checks(surface, limit, flip=False):
    """The one-sided checks for a surface, as limits on the quantity being kriged."""
    b = bounds[bounds.surface == surface].copy()
    b["kind"] = b.kind.replace({"deeper than": "above"})
    if flip:
        b["kind"] = b.kind.map({"above": "below", "below": "above"})
    b["limit"] = limit(b)
    return b.reset_index(drop=True)




sup = dedupe(picks[picks.surface == "rockhead"])
xs_, ys_, v_ = sup.x.values, sup.y.values, sup.depth.values
vgm["Superficial thickness"] = fit_spherical(xs_, ys_, v_, LAGS)
show("Superficial thickness")
print(f"  {len(sup)} thicknesses (after averaging repeat holes): median {sup.depth.median():.2f} m, "
      f"max {sup.depth.max():.2f} m ({sup.LOCA_ID[sup.depth.idxmax()]})")
f_sup, _, ex_sup, _ = honour("Rockhead", xs_, ys_, v_, vgm["Superficial thickness"][0], checks("rockhead", lambda b: b.GL - b.z))
p, s = loocv(xs_, ys_, v_, vgm["Superficial thickness"][0], extra=ex_sup)
sup["pred"], sup["sd"] = np.clip(p, 0, None), s
cv["Superficial thickness"] = scores(sup.depth, sup.pred, sup.sd)
thick_sup, sd_sup = on_grid(f_sup)
thick_sup = np.clip(thick_sup, 0, None)
rockhead = ground - thick_sup


# ------------------------------------------------------------------ Step 5: base of the Inferior Oolite
section("Step 5: base of the Inferior Oolite: regional dip plus kriged residuals")
bio_all = dedupe(picks[picks.surface == "base_io"])
bio = bio_all[~bio_all.LOCA_ID.isin(CAMBER)].reset_index(drop=True)
res, coef = residuals(bio.x.values, bio.y.values, bio.z.values)
grad = np.hypot(coef[1], coef[2])
azimuth = (np.degrees(np.arctan2(-coef[1], -coef[2])) + 360) % 360      # direction the surface falls towards
print(f"{len(bio)} holes (RC507-RC509 held back). Fitted plane dips {np.degrees(np.arctan(grad)):.2f}° "
      f"({grad * 1000:.1f} m per km) towards {azimuth:.0f}°; residuals {res.min():+.1f} to {res.max():+.1f} m")
vgm["Base of Inferior Oolite"] = fit_spherical(bio.x.values, bio.y.values, res, LAGS)
show("Base of Inferior Oolite")
b_bio = checks("base_io", lambda b: b.z)
f_bio, coef, ex_bio, added_bio = honour("Base of Inferior Oolite", bio.x.values, bio.y.values, bio.z.values, vgm["Base of Inferior Oolite"][0],
                                        b_bio, trend=True)
print("  added:", ", ".join(added_bio) or "none")
p, s = loocv(bio.x.values, bio.y.values, bio.z.values, vgm["Base of Inferior Oolite"][0], True, ex_bio)
bio["pred"], bio["sd"] = p, s
cv["Base of Inferior Oolite"] = scores(bio.z, bio.pred, bio.sd)
base_io, sd_bio = on_grid(f_bio)


# The three low holes, against a surface built without them.
cam = bio_all[bio_all.LOCA_ID.isin(CAMBER)].reset_index(drop=True)
cam["regional"], cam["sd"] = f_bio(cam.x.values, cam.y.values)
cam["offset"] = cam.z - cam.regional
print("RC507-RC509 against the surface built from the other holes:")
print(cam[["LOCA_ID", "z", "regional", "sd", "offset"]].round(2).to_string(index=False))
# ...and what happens to everyone else if they are simply kriged in with the rest.
res_w, _ = residuals(bio_all.x.values, bio_all.y.values, bio_all.z.values)
vgm_with = fit_spherical(bio_all.x.values, bio_all.y.values, res_w, LAGS)[0]
f_with, _ = krige(np.r_[bio_all.x, ex_bio[0]], np.r_[bio_all.y, ex_bio[1]], np.r_[bio_all.z, ex_bio[2]], vgm_with,
                  trend=True)
p, s = loocv(bio_all.x.values, bio_all.y.values, bio_all.z.values, vgm_with, True, ex_bio)
others = ~bio_all.LOCA_ID.isin(CAMBER)
cv_with = scores(bio_all.z[others], p[others], s[others])
print(f"kriging them in with the rest instead: LOOCV RMSE at the other {others.sum()} holes "
      f"{cv_with['RMSE']:.2f} m (without them {cv['Base of Inferior Oolite']['RMSE']:.2f} m); the variogram "
      f"nugget rises from {vgm['Base of Inferior Oolite'][0]['nugget']:.1f} to {vgm_with['nugget']:.1f} m²")
base_io_with, _ = on_grid(f_with)
base_io_regional = base_io.copy()


# Deliberate handling: a block bounded by the nearest-hole rule. Each cell nearer to one of the three
# than to any other hole that reached the contact, and within 60 m of it, drops by that hole's offset.
d_good, _ = cKDTree(bio[["x", "y"]]).query(np.c_[gx.ravel(), gy.ravel()])
d_cam, i_cam = cKDTree(cam[["x", "y"]]).query(np.c_[gx.ravel(), gy.ravel()])
block = ((d_cam < d_good) & (d_cam <= 60)).reshape(ny, nx) & inside
base_io[block] += cam.offset.to_numpy()[i_cam.reshape(ny, nx)[block]]
print(f"low block: {block.sum()} cells ({block.sum() * CELL**2 / 1e4:.1f} ha) lowered by "
      f"{-cam.offset.max():.1f} to {-cam.offset.min():.1f} m")


# ------------------------------------------------------------------ Step 6: hang the other surfaces from it
section("Step 6: top of the Inferior Oolite and base of the Bridport Sand from thicknesses")
tio = dedupe(picks[picks.surface == "top_io"])
both = tio.LOCA_ID.isin(bio_all.LOCA_ID)
tio["base_here"] = np.where(both, tio.LOCA_ID.map(bio_all.set_index("LOCA_ID").z), at(base_io, tio.E, tio.N))
tio["thick"] = tio.z - tio.base_here
print("Inferior Oolite thickness = top seen in the log minus the base (seen if the hole goes through, else modelled)")
print(tio[["LOCA_ID", "z", "base_here", "thick"]].assign(through=both).round(2).to_string(index=False))
vgm["Inferior Oolite thickness"] = fit_spherical(tio.x.values, tio.y.values, tio.thick.values,
                                                 np.arange(0, 2050, 250.0))
show("Inferior Oolite thickness")
f_tio, _, ex_tio, added_tio = honour("Top of Inferior Oolite", tio.x.values, tio.y.values, tio.thick.values, vgm["Inferior Oolite thickness"][0],
                                     checks("top_io", lambda b: b.z - at(base_io, b.E, b.N)))
p, s = loocv(tio.x.values, tio.y.values, tio.thick.values, vgm["Inferior Oolite thickness"][0], extra=ex_tio)
tio["pred"], tio["sd"] = p + tio.base_here, s
cv["Top of Inferior Oolite"] = scores(tio.z, tio.pred, tio.sd)
thick_io, sd_tio = on_grid(f_tio)
top_io = base_io + thick_io
sd_tio = np.hypot(sd_tio, sd_bio)


bds = dedupe(picks[picks.surface == "base_bds"])
bds["z_top"] = bds.LOCA_ID.map(bio_all.set_index("LOCA_ID").z)
bds["thick"] = bds.z_top - bds.z
print(f"\nBridport Sand, {len(bds)} complete penetrations: " +
      ", ".join(f"{a} {b:.1f} m" for a, b in zip(bds.LOCA_ID, bds.thick)))
# Five thicknesses cannot give a variogram of their own: borrow the range of the base of the Inferior Oolite
# and use their variance as the sill.
vgm["Bridport Sand thickness"] = (dict(psill=bds.thick.var(), range=vgm["Base of Inferior Oolite"][0]["range"],
                                       nugget=0.0), experimental(bds.x.values, bds.y.values, bds.thick.values, LAGS))
show("Bridport Sand thickness")
f_bds, _, ex_bds, added_bds = honour("Base of Bridport Sand", bds.x.values, bds.y.values, bds.thick.values, vgm["Bridport Sand thickness"][0],
                                     checks("base_bds", lambda b: at(base_io, b.E, b.N) - b.z, flip=True))
print("  added (holes that end in Bridport Sand below the first model):", ", ".join(added_bds) or "none")
p, s = loocv(bds.x.values, bds.y.values, bds.thick.values, vgm["Bridport Sand thickness"][0], extra=ex_bds)
bds["pred"], bds["sd"] = bds.z_top - p, s
cv["Base of Bridport Sand"] = scores(bds.z, bds.pred, bds.sd)
thick_bds, sd_t = on_grid(f_bds)
base_bds = base_io - np.clip(thick_bds, 0, None)
sd_bds = np.hypot(sd_bio, sd_t)
print(f"modelled Bridport Sand thickness {np.nanmin(thick_bds[inside]):.1f} to {np.nanmax(thick_bds[inside]):.1f} m")
# ------------------------------------------------------------------ Step 7: the checks
section("Step 7: does the model hold together?")
print("leave-one-out cross-validation (prediction minus log, m; standardised errors should average 0, variance 1):")
print(pd.DataFrame(cv).T.round(2).to_string())


sep = (top_io - base_io)[inside]
sep2 = (base_io - base_bds)[inside]
print(f"\nno-crossing: top minus base of Inferior Oolite {np.nanmin(sep):.1f} to {np.nanmax(sep):.1f} m "
      f"({(sep < 0).sum()} cells cross); base of Inferior Oolite minus base of Bridport Sand {np.nanmin(sep2):.1f} "
      f"to {np.nanmax(sep2):.1f} m; rockhead minus ground {np.nanmax((rockhead - ground)[inside]):.1f} m at most")
raw = {"rockhead": rockhead, "top_io": top_io, "base_io": base_io, "base_bds": base_bds}
r = ((Y1 - bounds.N) // CELL).astype(int); c = ((bounds.E - X0) // CELL).astype(int)
bounds["model"] = [raw[s][i, j] for s, i, j in zip(bounds.surface, r, c)]
bounds["margin"] = np.where(bounds.kind == "below", bounds.z - bounds.model, bounds.model - bounds.z)
bounds.loc[bounds.kind == "deeper than", "margin"] = (bounds.z - bounds.model)[bounds.kind == "deeper than"]
bounds["ok"] = bounds.margin >= -0.5                    # half a metre for the 10 m cells and logged levels
tab = bounds.groupby(["surface", "kind"]).agg(holes=("ok", "size"), agree=("ok", "sum"),
                                              worst_m=("margin", "min")).reindex(SURFACES, level=0)
print("\nholes that miss a surface: does the model agree with them?")
print(tab.round(1).to_string())
print("disagreements:")
print(bounds[~bounds.ok][["LOCA_ID", "surface", "kind", "z", "model", "margin"]].round(1).to_string(index=False))


# Stack the surfaces: where a younger surface is lower, the older one has been eroded and takes its level.
stack = {"ground": ground.copy(), "rockhead": rockhead}
stack["top_io"] = np.fmin(top_io, stack["rockhead"])
stack["base_io"] = np.fmin(base_io, stack["top_io"])
stack["base_bds"] = np.fmin(base_bds, stack["base_io"])
for g in stack.values():
    g[~inside] = np.nan
MODEL_BASE = float(np.floor((np.nanmin(stack["base_bds"]) - 10) / 10) * 10)
present = {u: np.nanmean((stack[a] - stack[b])[inside] > 0.01) for u, a, b in
           [("SUP", "ground", "rockhead"), ("GOG", "rockhead", "top_io"), ("IOG", "top_io", "base_io"),
            ("BDS", "base_io", "base_bds")]}
present["LIAS"] = 1.0
print(f"\nmodel base {MODEL_BASE:.0f} m OD. Share of the model footprint where each unit is present: " +
      ", ".join(f"{UNIT.name[u].split(' (')[0]} {v:.0%}" for u, v in present.items()))


# ------------------------------------------------------------------ Step 8: write the model
section("Step 8: write the surfaces")
FILES = {"ground": "ground_mOD", "rockhead": "rockhead_mOD", "top_io": "top_inferior_oolite_mOD",
         "base_io": "base_inferior_oolite_mOD", "base_bds": "base_bridport_sand_mOD"}
sds = {"rockhead": sd_sup, "top_io": sd_tio, "base_io": sd_bio, "base_bds": sd_bds}




def write_tif(name, a):
    with rasterio.open(MODEL / f"{name}.tif", "w", driver="GTiff", width=nx, height=ny, count=1, dtype="float32",
                       crs="EPSG:27700", transform=transform, nodata=np.nan, compress="deflate") as dst:
        dst.write(np.where(inside, a, np.nan).astype("float32"), 1)




for s, name in FILES.items():
    write_tif(name, stack[s])
for s, a in sds.items():
    write_tif(f"sd_{FILES[s]}", a)
write_tif("low_block", block.astype("float32"))
tops = list(FILES.values())
meta = {"model_base_mOD": MODEL_BASE, "cell_m": CELL, "crs": "EPSG:27700",
        "units": [{"code": u, "name": UNIT.name[u], "colour": UNIT.colour[u], "top": tops[i],
                   "base": tops[i + 1] if i + 1 < len(tops) else None} for i, u in enumerate(ORDER)],
        "credit": CREDIT}
(MODEL / "model.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")
pd.concat([sup.assign(surface="superficial thickness", obs=sup.depth), bio.assign(obs=bio.z),
           tio.assign(obs=tio.z), bds.assign(obs=bds.z)])[["LOCA_ID", "surface", "E", "N", "obs", "pred", "sd"]] \
    .round(2).to_csv(OUT / "loocv.csv", index=False)
print("wrote", ", ".join(sorted(p.name for p in MODEL.glob("*"))), "and loocv.csv")


# ------------------------------------------------------------------ Step 9: figures
section("Step 9: figures")
EXT = [X0, X1, Y0, Y1]
hill = LightSource(azdeg=315, altdeg=40).hillshade(np.nan_to_num(dtm5, nan=np.nanmean(dtm5)), vert_exag=2,
                                                    dx=RES, dy=RES)
# The route corridor, vale to plateau, for the cross-section.
LINE = np.array([(392_050, 215_640), (392_650, 215_880), (393_080, 215_870), (393_330, 216_030),
                 (393_700, 216_000), (394_000, 215_770), (394_240, 215_150), (394_680, 213_850)], float)
seg = np.diff(LINE, axis=0)
seglen = np.hypot(*seg.T)
chain0 = np.r_[0, np.cumsum(seglen)]




def project(E, N):
    """Chainage along LINE and distance off it for each point."""
    P = np.c_[E, N][:, None, :]
    t = np.clip(((P - LINE[:-1]) * seg).sum(-1) / seglen**2, 0, 1)
    foot = LINE[:-1] + t[..., None] * seg
    d = np.hypot(*(P - foot).transpose(2, 0, 1))
    k = d.argmin(axis=1)
    return chain0[k] + t[np.arange(len(k)), k] * seglen[k], d[np.arange(len(k)), k]




def frame(ax, ext=None, grid=True):
    ext = ext or EXT
    ax.set_xlim(ext[0], ext[1]); ax.set_ylim(ext[2], ext[3])
    ax.set_aspect("equal")
    ax.ticklabel_format(useOffset=False, style="plain")
    ax.tick_params(labelsize=6.5)
    for lab in ax.get_xticklabels() + ax.get_yticklabels():
        lab.set_rotation(0)




def footprint(ax, colour="#333333"):
    ax.contour(gx + X0, gy + Y0, inside.astype(float), levels=[0.5], colors=colour, linewidths=0.8)




role = pd.Series("logged, rockhead only", index=holes.LOCA_ID)
role[bounds.loc[bounds.kind == "deeper than", "LOCA_ID"]] = "ends in superficial deposits"
role[picks.loc[picks.surface == "top_io", "LOCA_ID"]] = "top of Inferior Oolite seen"
role[picks.loc[picks.surface == "base_io", "LOCA_ID"]] = "base of Inferior Oolite seen"
role[CAMBER] = "RC507-RC509 (low block)"
STYLE = {"logged, rockhead only": dict(marker="o", s=10, c="#ffffff", edgecolors="#333333", lw=0.6),
         "ends in superficial deposits": dict(marker="o", s=10, c="#c98c5a", edgecolors="#333333", lw=0.6),
         "top of Inferior Oolite seen": dict(marker="^", s=30, c="#f4cd72", edgecolors="black", lw=0.7),
         "base of Inferior Oolite seen": dict(marker="s", s=24, c="#1f3b73", edgecolors="white", lw=0.5),
         "RC507-RC509 (low block)": dict(marker="s", s=30, c="#b2182b", edgecolors="white", lw=0.5)}


# Figure 2: the data.
fig, ax = plt.subplots(figsize=(8.6, 8.2))
ax.imshow(hill, cmap="gray", extent=EXT, vmin=0, vmax=1.15)
footprint(ax)
ax.plot(*LINE.T, color="#b2182b", lw=1.4, ls="--")
ax.annotate("section, figure 7", LINE[0], xytext=(-8, -16), textcoords="offset points", fontsize=7, color="#b2182b")
for k, st in STYLE.items():
    sel = holes[holes.LOCA_ID.map(role) == k]
    ax.scatter(sel.LOCA_NATE, sel.LOCA_NATN, label=f"{k} ({len(sel)})", zorder=3, **st)
for lid in ["CP217", "OH413", "DSRC404", "OH411", "DSRCOH400", "DSRCOH306"]:
    h = holes.set_index("LOCA_ID").loc[lid]
    ax.annotate(lid, (h.LOCA_NATE, h.LOCA_NATN), xytext=(5, 3), textcoords="offset points", fontsize=7,
                bbox=dict(boxstyle="round,pad=0.15", fc="white", ec="none", alpha=0.8))
ax.plot([], [], color="#333333", lw=0.8, label=f"model footprint (within {REACH:.0f} m of a hole)")
ax.legend(loc="lower left", fontsize=7, framealpha=0.9)
frame(ax)
ax.set_title(f"{len(holes)} logged holes and what each one tells the model", fontsize=10)
fig.text(0.01, 0.01, f"{CREDIT} Hillshade from LIDAR DTM. British National Grid (m).", fontsize=6.5,
         color="#444444")
fig.savefig(OUT / "02-borehole-data-map.png")
plt.close(fig)


# Figure 3: variograms and cross-validation.
names = list(vgm)
cvd = {"Superficial thickness": (sup.depth, sup.pred, "superficial thickness (m)"),
       "Base of Inferior Oolite": (bio.z, bio.pred, "base of Inferior Oolite (m OD)"),
       "Inferior Oolite thickness": (tio.z, tio.pred, "top of Inferior Oolite (m OD)"),
       "Bridport Sand thickness": (bds.z, bds.pred, "base of Bridport Sand (m OD)")}
cvname = dict(zip(names, cv))
fig, axes = plt.subplots(2, 4, figsize=(13, 6.4))
for i, n in enumerate(names):
    p, (lag, gam, cnt) = vgm[n]
    ax = axes[0, i]
    ax.scatter(lag, gam, s=8 + 40 * cnt / cnt.max(), color="#4a6fa5", label="pairs of holes, binned")
    h = np.linspace(0, 2000, 200)
    ax.plot(h, spherical(h, p["psill"], p["range"], p["nugget"]), color="black", lw=1.2, label="spherical model")
    ax.set_xlim(0, 2000); ax.set_ylim(0, None)
    ax.set_title(("Residual of the base of the Inferior Oolite" if n.startswith("Base of") else n), fontsize=8.5)
    ax.set_xlabel("separation (m)", fontsize=7.5); ax.set_ylabel("semivariance (m²)", fontsize=7.5)
    ax.text(0.97, 0.05, f"range {p['range']:.0f} m\nnugget {p['nugget']:.1f} m²", transform=ax.transAxes,
            ha="right", fontsize=7)
    obs, pred, lab = cvd[n]
    sc = cv[cvname[n]]
    ax = axes[1, i]
    lo, hi = min(obs.min(), pred.min()), max(obs.max(), pred.max())
    pad = (hi - lo) * 0.06
    ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], color="#999999", lw=0.8)
    ax.scatter(obs, pred, s=14, color="#d98f3e", edgecolors="#5b3a12", lw=0.4)
    ax.set_xlim(lo - pad, hi + pad); ax.set_ylim(lo - pad, hi + pad); ax.set_aspect("equal")
    ax.set_xlabel(f"logged {lab}", fontsize=7.5); ax.set_ylabel("predicted with that hole left out", fontsize=7.5)
    ax.text(0.04, 0.96, f"n = {sc['n']}\nRMSE {sc['RMSE']:.1f} m\nmean error {round(sc['mean_error'], 1) + 0.0:+.1f} m\n"
            f"variance of standardised\nerrors {sc['var_SE']:.2f}", transform=ax.transAxes, va="top", fontsize=7)
    worst = np.argsort(-np.abs(np.asarray(pred) - np.asarray(obs)))[:2]
    lids = {"Superficial thickness": sup, "Base of Inferior Oolite": bio, "Inferior Oolite thickness": tio,
            "Bridport Sand thickness": bds}[n].LOCA_ID.values
    for w in worst:
        ax.annotate(lids[w], (np.asarray(obs)[w], np.asarray(pred)[w]), xytext=(4, -9), textcoords="offset points",
                    fontsize=6.5)
for ax in axes.ravel():
    ax.tick_params(labelsize=6.5)
    ax.spines[["top", "right"]].set_visible(False)
axes[0, 0].legend(fontsize=6.5, loc="upper left", frameon=False)
fig.suptitle("Top: how alike two holes are against how far apart they are. Bottom: leave-one-out cross-validation",
             fontsize=10, x=0.01, ha="left")
fig.text(0.01, -0.01, CREDIT, fontsize=6.5, color="#444444")
fig.tight_layout()
fig.savefig(OUT / "03-variograms-cross-validation.png")
plt.close(fig)


# Figure 4: surfaces and their kriging standard deviation.
rows, cols = np.where(inside)
ZEXT = [X0 + cols.min() * CELL - 100, X0 + (cols.max() + 1) * CELL + 100,
        Y1 - (rows.max() + 1) * CELL - 100, Y1 - rows.min() * CELL + 100]
crop = [Y1 - ZEXT[3], Y1 - ZEXT[2], ZEXT[0] - X0, ZEXT[1] - X0]
show_ = {"superficial thickness (m)": (thick_sup, sd_sup, "YlOrBr", None),
         "base of Inferior Oolite (m OD)": (base_io, sd_bio, "viridis", stack["base_io"] < base_io - 0.01),
         "top of Inferior Oolite (m OD)": (top_io, sd_tio, "viridis", stack["top_io"] < top_io - 0.01)}
fig, axes = plt.subplots(2, 3, figsize=(13, 9.6))
for i, (lab, (val, sd, cmap, eroded)) in enumerate(show_.items()):
    for j, (arr, cm, title) in enumerate([(val, cmap, lab), (sd, "magma_r", "kriging standard deviation (m)")]):
        ax = axes[j, i]
        ax.imshow(hill, cmap="gray", extent=EXT, vmin=0, vmax=1.3)
        a = np.where(inside, arr, np.nan)
        im = ax.imshow(a, cmap=cm, extent=EXT, alpha=0.85,
                       vmin=0 if j else np.nanpercentile(a, 1), vmax=np.nanpercentile(a, 99))
        if j == 0 and i > 0:
            cs = ax.contour(gx + X0, gy + Y0, a, levels=np.arange(150, 320, 10), colors="black", linewidths=0.5)
            ax.clabel(cs, fontsize=5.5, fmt="%d")
            ax.contourf(gx + X0, gy + Y0, np.where(inside, eroded, np.nan).astype(float), levels=[0.5, 1.5],
                        hatches=["////"], colors="none")
        if j == 0 and i == 1:
            ax.contour(gx + X0, gy + Y0, block.astype(float), levels=[0.5], colors="#b2182b", linewidths=1.2)
        footprint(ax)
        pts = {0: sup, 1: bio, 2: tio}[i]
        ax.scatter(pts.E, pts.N, s=5, c="black", zorder=3)
        frame(ax, ZEXT)
        cb = fig.colorbar(im, ax=ax, shrink=0.75, pad=0.02)
        cb.ax.tick_params(labelsize=6.5)
        ax.set_title(title if j else title + (" (hatched: eroded, above rockhead)" if i else ""), fontsize=8)
fig.suptitle("The kriged surfaces (top) and how uncertain each one is (bottom). Dots: holes that log the contact; "
             "red outline: the low block", fontsize=10, x=0.01, ha="left")
fig.text(0.01, 0.0, f"{CREDIT} Hillshade from LIDAR DTM. British National Grid (m).", fontsize=6.5, color="#444444")
fig.tight_layout()
fig.savefig(OUT / "04-surfaces-and-uncertainty.png")
plt.close(fig)


# Figure 5: the low block at RC507-RC509.
c0 = cam[["E", "N"]].mean().to_numpy()
ZB = [c0[0] - 280, c0[0] + 280, c0[1] - 200, c0[1] + 200]
fig = plt.figure(figsize=(13, 5.2))
ax = fig.add_axes([0.03, 0.1, 0.36, 0.8])
ax.imshow(hill, cmap="gray", extent=EXT, vmin=0, vmax=1.3)
im = ax.imshow(np.where(inside, base_io, np.nan), cmap="viridis", extent=EXT, alpha=0.8, vmin=185, vmax=235)
cs = ax.contour(gx + X0, gy + Y0, np.where(inside, base_io, np.nan), levels=np.arange(180, 240, 2.5),
                colors="black", linewidths=0.4)
ax.contour(gx + X0, gy + Y0, block.astype(float), levels=[0.5], colors="#b2182b", linewidths=1.4)
near = bio_all[(bio_all.E.between(*ZB[:2])) & (bio_all.N.between(*ZB[2:]))]
for r in near.itertuples():
    lowb = r.LOCA_ID in CAMBER
    ax.scatter(r.E, r.N, s=26, marker="s", c="#b2182b" if lowb else "#1f3b73", edgecolors="white", lw=0.5, zorder=3)
    ax.annotate(f"{r.LOCA_ID} {r.z:.0f}" if lowb else f"{r.z:.0f}", (r.E, r.N),
                xytext={"RC507": (-14, -14), "RC508": (6, -12), "RC509": (6, 5)}.get(
                    r.LOCA_ID, (-18, 6) if np.hypot(cam.E - r.E, cam.N - r.N).min() < 40 else (4, 4)),
                textcoords="offset points", fontsize=6.5, color="#b2182b" if lowb else "black",
                bbox=dict(fc="white", ec="none", alpha=0.75, pad=0.5))
frame(ax, ZB)
fig.colorbar(im, ax=ax, shrink=0.8, pad=0.02, label="base of Inferior Oolite (m OD)").ax.tick_params(labelsize=6.5)
ax.set_title("Final surface; red outline: the low block. Labels: logged level (m OD)", fontsize=8)
# profile through the three holes
u = np.r_[np.cos(np.radians(90 - azimuth)), np.sin(np.radians(90 - azimuth))]
ends = np.array([c0 - 280 * u, c0 + 280 * u])
t = np.linspace(0, 560, 113)
pe, pn = ends[0][0] + t * u[0], ends[0][1] + t * u[1]
ax.plot(*ends.T, color="white", lw=2.2); ax.plot(*ends.T, color="black", lw=1, ls="--")
ax2 = fig.add_axes([0.47, 0.13, 0.3, 0.74])
for g, lab, st in [(base_io_regional, "built without RC507-RC509", dict(color="#1f3b73", lw=1.4)),
                   (base_io_with, "all holes kriged together", dict(color="#7a7a7a", lw=1.2, ls=":")),
                   (base_io, "final: regional surface plus the low block", dict(color="#b2182b", lw=1.6))]:
    ax2.plot(t, at(g, pe, pn), label=lab, **st)
sd_line = at(sd_bio, pe, pn)
ax2.fill_between(t, at(base_io_regional, pe, pn) - sd_line, at(base_io_regional, pe, pn) + sd_line,
                 color="#1f3b73", alpha=0.12, lw=0, label="± one kriging standard deviation")
off = (np.c_[bio_all.E, bio_all.N] - ends[0]) @ np.c_[-u[1], u[0]].T
along = (np.c_[bio_all.E, bio_all.N] - ends[0]) @ u
sel = (np.abs(off.ravel()) < 120) & (along >= 0) & (along <= 560)
for r, a_ in zip(bio_all[sel].itertuples(), along[sel]):
    lowb = r.LOCA_ID in CAMBER
    ax2.scatter(a_, r.z, marker="s", s=28, c="#b2182b" if lowb else "#1f3b73", edgecolors="white", zorder=3)
    if lowb:
        ax2.annotate(r.LOCA_ID, (a_, r.z), xytext=(5, -3), textcoords="offset points", fontsize=6.5, color="#b2182b")
ax2.set_xlabel(f"distance along the profile, towards {azimuth:.0f}° (m); holes within 120 m", fontsize=7.5)
ax2.set_ylabel("base of Inferior Oolite (m OD)", fontsize=7.5)
ax2.legend(fontsize=6.5, loc="lower left", frameon=False)
ax2.tick_params(labelsize=6.5); ax2.spines[["top", "right"]].set_visible(False)
ax2.set_title("Profile down the regional dip, through the three holes", fontsize=8)
ax3 = fig.add_axes([0.83, 0.13, 0.15, 0.74])
vals = [cv["Base of Inferior Oolite"]["RMSE"], cv_with["RMSE"]]
ax3.bar([0, 1], vals, color=["#1f3b73", "#9a9a9a"], width=0.6)
for i, v in enumerate(vals):
    ax3.text(i, v + 0.1, f"{v:.2f} m", ha="center", fontsize=7.5)
ax3.set_xticks([0, 1], ["held out,\nthen a block", "kriged in\nwith the rest"], fontsize=7)
ax3.set_ylabel(f"LOOCV RMSE at the other {others.sum()} holes (m)", fontsize=7.5)
ax3.tick_params(labelsize=6.5); ax3.spines[["top", "right"]].set_visible(False)
ax3.set_title("Error elsewhere", fontsize=8)
fig.text(0.01, 0.0, f"{CREDIT} Hillshade from LIDAR DTM. British National Grid (m).", fontsize=6.5, color="#444444")
fig.savefig(OUT / "05-low-block.png")
plt.close(fig)


# Figure 6: thicknesses, crossing and the one-sided checks.
fig = plt.figure(figsize=(13, 5.4))
ax = fig.add_axes([0.02, 0.1, 0.34, 0.82])
ax.imshow(hill, cmap="gray", extent=EXT, vmin=0, vmax=1.3)
im = ax.imshow(np.where(inside, thick_io, np.nan), cmap="cividis", extent=EXT, alpha=0.85)
footprint(ax)
ax.scatter(tio.E, tio.N, s=18, marker="^", c="#f4cd72", edgecolors="black", lw=0.6, zorder=3)
for r in tio[tio.LOCA_ID.isin(bio_all.LOCA_ID)].itertuples():
    ax.annotate(f"{r.thick:.0f}", (r.E, r.N), xytext=(4, 2), textcoords="offset points", fontsize=6.5,
                bbox=dict(fc="white", ec="none", alpha=0.75, pad=0.4))
frame(ax, ZEXT)
fig.colorbar(im, ax=ax, shrink=0.8, pad=0.02, label="Inferior Oolite thickness (m)").ax.tick_params(labelsize=6.5)
ax.set_title("Kriged Inferior Oolite thickness; labels: holes that go right through", fontsize=8)
ax2 = fig.add_axes([0.45, 0.13, 0.22, 0.75])
ax2.hist(sep, bins=40, color="#d98f3e", alpha=0.85, label="top minus base of Inferior Oolite")
ax2.hist(sep2, bins=40, color="#b8a96a", alpha=0.85, label="base of Inferior Oolite minus\nbase of Bridport Sand")
ax2.axvline(0, color="#b2182b", lw=1)
ax2.text(1, ax2.get_ylim()[1] * 0.97, "crossing\nwould show\nleft of here", fontsize=6.5, color="#b2182b", va="top")
ax2.set_xlim(-8, None)
ax2.set_xlabel("vertical gap between surfaces, every model cell (m)", fontsize=7.5)
ax2.set_ylabel("cells", fontsize=7.5)
ax2.legend(fontsize=6.5, frameon=False, loc="upper right")
ax2.tick_params(labelsize=6.5); ax2.spines[["top", "right"]].set_visible(False)
ax2.set_title(f"No crossing: {(sep < 0).sum() + (sep2 < 0).sum()} of {2 * len(sep):,} cell gaps below zero",
              fontsize=8)
ax3 = fig.add_axes([0.75, 0.13, 0.23, 0.75])
sumr = bounds.groupby("surface").agg(n=("ok", "size"), ok=("ok", "sum")).reindex(SURFACES)
first = {"rockhead": passes["Rockhead"], "top_io": passes["Top of Inferior Oolite"],
         "base_io": passes["Base of Inferior Oolite"], "base_bds": passes["Base of Bridport Sand"]}
yy = np.arange(len(sumr))[::-1]
ax3.barh(yy + 0.2, [first[s][0] - first[s][1] for s in sumr.index], 0.38, color="#bdbdbd",
         label="agree with the first model")
ax3.barh(yy - 0.2, sumr.ok, 0.38, color="#1f3b73", label="agree with the final model (0.5 m slack)")
for y_, s in zip(yy, sumr.index):
    ax3.text(first[s][0] - first[s][1] + 1, y_ + 0.2, f"{first[s][0] - first[s][1]} of {first[s][0]}",
             fontsize=6.5, va="center")
    ax3.text(sumr.ok[s] + 1, y_ - 0.2, f"{sumr.ok[s]} of {sumr.n[s]}", fontsize=6.5, va="center")
ax3.set_yticks(yy, ["Rockhead", "Top of Inferior\nOolite", "Base of Inferior\nOolite", "Base of Bridport\nSand"],
               fontsize=7)
ax3.set_xlabel("holes that stop short of, or start below, the surface", fontsize=7.5)
ax3.legend(fontsize=6.5, frameon=False, loc="upper right")
ax3.tick_params(labelsize=6.5); ax3.spines[["top", "right"]].set_visible(False)
ax3.set_xlim(0, sumr.n.max() * 1.25)
ax3.set_title("One-sided checks", fontsize=8)
fig.text(0.01, 0.0, f"{CREDIT} Hillshade from LIDAR DTM. British National Grid (m).", fontsize=6.5, color="#444444")
fig.savefig(OUT / "06-thickness-and-checks.png")
plt.close(fig)


# Figure 7: the section.
STEP = 10.0
ch = np.arange(0, chain0[-1], STEP)
k_ = np.clip(np.searchsorted(chain0, ch, side="right") - 1, 0, len(seg) - 1)
fr = (ch - chain0[k_]) / seglen[k_]
se, sn = LINE[k_, 0] + fr * seg[k_, 0], LINE[k_, 1] + fr * seg[k_, 1]
prof = {s: at(stack[s], se, sn) for s in stack}
prof["model_base"] = np.where(np.isfinite(prof["ground"]), MODEL_BASE, np.nan)
SEQ = [("SUP", "ground", "rockhead"), ("GOG", "rockhead", "top_io"), ("IOG", "top_io", "base_io"),
       ("BDS", "base_io", "base_bds"), ("LIAS", "base_bds", "model_base")]
VS = 5.0
fig, ax = plt.subplots(figsize=(14, 4.4))
for u_, a_, b_ in SEQ:
    ax.fill_between(ch, prof[b_], prof[a_], color=UNIT.colour[u_], lw=0, label=UNIT.name[u_])
ax.plot(ch, at(ground, se, sn), color="black", lw=0.9, label="LIDAR ground")
sdb = at(sd_bio, se, sn)
ax.plot(ch, prof["base_io"] - sdb, color="#1f3b73", lw=0.7, ls="--")
ax.plot(ch, prof["base_io"] + sdb, color="#1f3b73", lw=0.7, ls="--",
        label="base of Inferior Oolite ± one kriging SD")
cE, dE = project(holes.LOCA_NATE.to_numpy(), holes.LOCA_NATN.to_numpy())
near = holes.assign(ch=cE, off=dE)[dE <= 75]
w = 14
for h in near.itertuples():
    for r in runs[runs.LOCA_ID == h.LOCA_ID].itertuples():
        ax.add_patch(plt.Rectangle((h.ch - w / 2, r.LOCA_GL - r.base), w, r.base - r.top, fc=UNIT.colour[r.group],
                                   ec="black", lw=0.5, zorder=4))
for lid in ["CP217", "RC508", "DSRCOH306", "OH411", "OH413", "DSRCOH400"]:
    if lid in set(near.LOCA_ID):
        h = near.set_index("LOCA_ID").loc[lid]
        ax.annotate(lid, (h.ch, h.LOCA_GL + 4), ha="center", fontsize=6.5, rotation=90, va="bottom")
ax.set_xlim(0, ch[-1]); ax.set_ylim(MODEL_BASE, np.nanmax(prof["ground"]) + 75)   # headroom so the legend clears the hole labels
ax.set_aspect(VS)
ax.set_xlabel("distance along the section (m)", fontsize=8); ax.set_ylabel("m OD", fontsize=8)
ax.tick_params(labelsize=7)
for cx_ in chain0[1:-1]:
    ax.axvline(cx_, color="#555555", lw=0.4, ls=":")
ax.legend(fontsize=7, ncol=4, loc="upper right", frameon=True, framealpha=0.9)
ax.set_title(f"Section along the route corridor, vale (west) to plateau (south-east), {VS:.0f}x vertical "
             f"exaggeration. Logs within 75 m, collapsed to model units", fontsize=9, loc="left")
fig.text(0.01, 0.0, f"{CREDIT} Dotted lines: bends in the section.", fontsize=6.5, color="#444444")
fig.savefig(OUT / "07-section.png")
plt.close(fig)


# Figure 1: the model in 3D, with the superficial deposits lifted off.
X3, Y3 = np.stack([gx, gx], -1), np.stack([gy, gy], -1)
lower = {"SUP": "rockhead", "GOG": "top_io", "IOG": "base_io", "BDS": "base_bds", "LIAS": None}
upper = {"SUP": "ground", "GOG": "rockhead", "IOG": "top_io", "BDS": "base_io", "LIAS": "base_bds"}
pl = pv.Plotter(off_screen=True, window_size=(1800, 1150))
pl.set_background("white")
for u_ in ORDER:
    top = np.nan_to_num(stack[upper[u_]], nan=MODEL_BASE)
    base = np.full_like(top, MODEL_BASE) if lower[u_] is None else np.nan_to_num(stack[lower[u_]], nan=MODEL_BASE)
    g = pv.StructuredGrid(X3, Y3, np.stack([base, top], -1) * VE)
    g.point_data["keep"] = np.stack([inside & (top - base > 0.05)] * 2, -1).ravel(order="F").astype(float)
    solid = g.threshold(0.5, scalars="keep", all_scalars=True)
    if solid.n_cells == 0 or u_ == "SUP":
        continue
    pl.add_mesh(solid.extract_surface(algorithm=None), color=UNIT.colour[u_], smooth_shading=True, split_sharp_edges=True,
                ambient=0.35, diffuse=0.7, specular=0.0)
gnd = pv.StructuredGrid(gx, gy, np.nan_to_num(stack["ground"], nan=MODEL_BASE) * VE + 2)
gnd.point_data["keep"] = inside.ravel(order="F").astype(float)
pl.add_mesh(gnd.threshold(0.5, scalars="keep", all_scalars=True).extract_surface(algorithm=None), color="#c98c5a",
            opacity=0.18)
pins = pv.merge([pv.Line((e - X0, n - Y0, z * VE), (e - X0, n - Y0, (z + 18) * VE))
                 for e, n, z in zip(holes.LOCA_NATE, holes.LOCA_NATN, holes.LOCA_GL)])
pl.add_mesh(pins, color="black", line_width=2)
pl.camera_position = [(-1500, -2600, 2400 * VE / 3 + 900), (2100, 2300, 160 * VE), (0, 0, 1)]
pl.camera.zoom(1.35)
img = pl.screenshot(return_img=True)
pl.close()




def compose(img, units, path, title, note=""):
    """Put a PyVista screenshot beside a proper legend (PyVista's own legend tints the text)."""
    ink = np.where((img[..., :3] < 245).any(axis=2))                 # trim the empty white margins
    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")
    if note:
        fig.text(0.01, 0.005, note, fontsize=7, color="#444444")
    handles = [Patch(facecolor=UNIT.colour[u], edgecolor="#555555", lw=0.4, label=UNIT.name[u]) for u in units]
    handles.append(Patch(facecolor="#c98c5a", alpha=0.25, edgecolor="#555555", lw=0.4,
                         label="LIDAR ground (superficial\ndeposits lifted off)"))
    handles.append(Line2D([], [], color="black", lw=1.2, label="logged hole"))
    fig.legend(handles=handles, loc="center left", bbox_to_anchor=(0.79, 0.5), fontsize=7.5, frameon=False,
               title="Model unit", title_fontsize=8, alignment="left")
    fig.savefig(path, dpi=150)
    plt.close(fig)




compose(img, ["GOG", "IOG", "BDS", "LIAS"], OUT / "01-ground-model-3d.png",
        "A417 ground model: kriged stratum surfaces under LIDAR ground, looking north-east up Crickley Hill",
        f"{CREDIT} {VE:.0f}x vertical exaggeration; model base {MODEL_BASE:.0f} m OD.")
print("wrote", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))