Which rung on the k-mesh ladder does this crystal need?¶

k_points.k_index.qrf — the quantile forest published as 4050a-aas85.

This model does not return a k-point spacing. It returns a rung: a position on the ordered ladder of meshes that this particular crystal can have, where rung 1 is the Γ-only (1, 1, 1) mesh. The same rung is a different mesh for a different cell, so the ladder has to be rebuilt per structure before the number means anything.

This notebook builds that ladder, checks it against the published labels, and then shows where the model works and where it does not.

Setup¶

Unlike QRF95 this model depends on nothing else: its 174 inputs are all deterministic functions of the crystal.

In [1]:

Copied!





import shutil
import warnings
from pathlib import Path

from goldilocks_ml.inference import load_model

# Several MC3D CIFs round coordinates on parse; the notice is not about this model.
warnings.filterwarnings("ignore", message="Issues encountered while parsing CIF")

work = Path("goldilocks-kindex")
work.mkdir(exist_ok=True)

ARTIFACTS = Path("../../local_runs/kindex-qrf-52713-55d86/model")
DEPOSIT = Path("../../deposits/k_points/k_index/qrf/model.json")
for name in ("k_index_qrf.pkl", "calibration.json"):
    if not (work / name).exists():
        shutil.copy(ARTIFACTS / name, work / name)
shutil.copy(DEPOSIT, work / "model.json")

model = load_model(work)
print(model.model_id)
print("decision rule:", model.record["decision"]["rule"])
print("bands (cut on the model's own median):")
for band in model.record["decision"]["bands"]:
    upper = (
        f"below rung {band['upper']:.0f}"
        if band["upper"] is not None
        else "rung 12 and above"
    )
    print(f"  {upper}: publishes q{band['level']}")
import shutil
import warnings
from pathlib import Path

from goldilocks_ml.inference import load_model

# Several MC3D CIFs round coordinates on parse; the notice is not about this model.
warnings.filterwarnings("ignore", message="Issues encountered while parsing CIF")

work = Path("goldilocks-kindex")
work.mkdir(exist_ok=True)

ARTIFACTS = Path("../../local_runs/kindex-qrf-52713-55d86/model")
DEPOSIT = Path("../../deposits/k_points/k_index/qrf/model.json")
for name in ("k_index_qrf.pkl", "calibration.json"):
    if not (work / name).exists():
        shutil.copy(ARTIFACTS / name, work / name)
shutil.copy(DEPOSIT, work / "model.json")

model = load_model(work)
print(model.model_id)
print("decision rule:", model.record["decision"]["rule"])
print("bands (cut on the model's own median):")
for band in model.record["decision"]["bands"]:
    upper = (
        f"below rung {band['upper']:.0f}"
        if band["upper"] is not None
        else "rung 12 and above"
    )
    print(f"  {upper}: publishes q{band['level']}")

k_points.k_index.qrf@cslr.v1
decision rule: quantile_by_band
bands (cut on the model's own median):
  below rung 7: publishes q0.9
  below rung 12: publishes q0.975
  rung 12 and above: publishes q0.975

The ladder¶

A mesh changes when ceil(|b_i| / k_distance) steps up, which happens exactly at k_distance = |b_i| / n. Those quotients are the only places the ladder can move, so probing between them enumerates it. This is what Goldilocks Core does to turn a rung into a mesh.

In [2]:

Copied!





import math

from pymatgen.core import Lattice, Structure


def mesh_at(structure, k_distance):
    b = structure.lattice.reciprocal_lattice
    return tuple(
        max(1, math.ceil(round(length / k_distance, 5))) for length in (b.a, b.b, b.c)
    )


def ladder(structure, min_k_distance=0.03):
    b = structure.lattice.reciprocal_lattice
    lengths = (b.a, b.b, b.c)
    changes = sorted(
        {
            round(length / n, 8)
            for length in lengths
            for n in range(1, math.floor(length / min_k_distance) + 1)
        },
        reverse=True,
    )
    probes = [changes[0] + 1.0]
    probes += [0.5 * (hi + lo) for hi, lo in zip(changes[:-1], changes[1:])]

    rungs, seen, previous = [], set(), None
    for probe in probes:
        current = mesh_at(structure, probe)
        # A jump of more than one k-point means a change point was never
        # enumerated, so the ladder has a hole from here on.
        if previous and any(now - was > 1 for was, now in zip(previous, current)):
            break
        previous = current
        if current not in seen:
            seen.add(current)
            rungs.append(current)
    return rungs


# rungs[0] is rung 1: this ladder is 1-based, so index it with rung - 1.
print(
    "rung 1 is always the Gamma-only mesh:",
    ladder(Structure(Lattice.cubic(4.0), ["Cu"], [[0, 0, 0]]))[0],
)
import math

from pymatgen.core import Lattice, Structure


def mesh_at(structure, k_distance):
    b = structure.lattice.reciprocal_lattice
    return tuple(
        max(1, math.ceil(round(length / k_distance, 5))) for length in (b.a, b.b, b.c)
    )


def ladder(structure, min_k_distance=0.03):
    b = structure.lattice.reciprocal_lattice
    lengths = (b.a, b.b, b.c)
    changes = sorted(
        {
            round(length / n, 8)
            for length in lengths
            for n in range(1, math.floor(length / min_k_distance) + 1)
        },
        reverse=True,
    )
    probes = [changes[0] + 1.0]
    probes += [0.5 * (hi + lo) for hi, lo in zip(changes[:-1], changes[1:])]

    rungs, seen, previous = [], set(), None
    for probe in probes:
        current = mesh_at(structure, probe)
        # A jump of more than one k-point means a change point was never
        # enumerated, so the ladder has a hole from here on.
        if previous and any(now - was > 1 for was, now in zip(previous, current)):
            break
        previous = current
        if current not in seen:
            seen.add(current)
            rungs.append(current)
    return rungs


# rungs[0] is rung 1: this ladder is 1-based, so index it with rung - 1.
print(
    "rung 1 is always the Gamma-only mesh:",
    ladder(Structure(Lattice.cubic(4.0), ["Cu"], [[0, 0, 0]]))[0],
)

rung 1 is always the Gamma-only mesh: (1, 1, 1)

Does it match the published labels?¶

The dataset record publishes both the rung and the mesh that rung was measured at. If this is the same ladder the labels were made on, rung n must reproduce the recorded mesh.

In [3]:

Copied!





import csv
import random
import re

SNAPSHOT = Path("../../local_data/snapshots/kindex-52713-55d86")
# The dataset record's own summary, which pairs every rung with the mesh it was
# measured at. Downloaded rather than committed: it is 700 kB of ground truth.
SUMMARY = work / "convergence_summary.csv"
if not SUMMARY.exists():
    from urllib.error import URLError
    from urllib.request import urlretrieve

    try:
        urlretrieve(
            "https://data-collections.psdi.ac.uk/api/records/52713-55d86"
            "/files/convergence_summary.csv/content",
            SUMMARY,
        )
    except (URLError, OSError):
        pass  # handled below: the check is skipped, not the whole notebook

rows = list(csv.reader(open(SNAPSHOT / "id_prop.csv")))
random.seed(0)
random.shuffle(rows)

if SUMMARY.exists():
    summary = {r["source_db_id"]: r for r in csv.DictReader(open(SUMMARY))}
    checked = exact = 0
    for sample_id, rung, _ in rows:
        if checked >= 30 or sample_id not in summary:
            continue
        structure = Structure.from_file(SNAPSHOT / f"{sample_id}.cif")
        rungs = ladder(structure)
        rung = int(rung)
        if rung > len(rungs):
            continue
        recorded = tuple(
            int(v) for v in re.findall(r"\d+", summary[sample_id]["k_mesh"])
        )
        checked += 1
        # Rung 1 is rungs[0]: a 1-based rung indexes this list at rung - 1.
        exact += rungs[rung - 1] == recorded
    print(f"{exact}/{checked} rungs reproduce the published mesh")
else:
    print("convergence_summary.csv not downloaded; skipping the check")
import csv
import random
import re

SNAPSHOT = Path("../../local_data/snapshots/kindex-52713-55d86")
# The dataset record's own summary, which pairs every rung with the mesh it was
# measured at. Downloaded rather than committed: it is 700 kB of ground truth.
SUMMARY = work / "convergence_summary.csv"
if not SUMMARY.exists():
    from urllib.error import URLError
    from urllib.request import urlretrieve

    try:
        urlretrieve(
            "https://data-collections.psdi.ac.uk/api/records/52713-55d86"
            "/files/convergence_summary.csv/content",
            SUMMARY,
        )
    except (URLError, OSError):
        pass  # handled below: the check is skipped, not the whole notebook

rows = list(csv.reader(open(SNAPSHOT / "id_prop.csv")))
random.seed(0)
random.shuffle(rows)

if SUMMARY.exists():
    summary = {r["source_db_id"]: r for r in csv.DictReader(open(SUMMARY))}
    checked = exact = 0
    for sample_id, rung, _ in rows:
        if checked >= 30 or sample_id not in summary:
            continue
        structure = Structure.from_file(SNAPSHOT / f"{sample_id}.cif")
        rungs = ladder(structure)
        rung = int(rung)
        if rung > len(rungs):
            continue
        recorded = tuple(
            int(v) for v in re.findall(r"\d+", summary[sample_id]["k_mesh"])
        )
        checked += 1
        # Rung 1 is rungs[0]: a 1-based rung indexes this list at rung - 1.
        exact += rungs[rung - 1] == recorded
    print(f"{exact}/{checked} rungs reproduce the published mesh")
else:
    print("convergence_summary.csv not downloaded; skipping the check")

30/30 rungs reproduce the published mesh

A crystal it handles well¶

The training set is MC3D: bulk crystals from experimental databases, mostly with large cells.

In [4]:

Copied!





clathrate = Structure.from_file(SNAPSHOT / "56721.cif")
rungs = ladder(clathrate)
predicted = int(model.predict(clathrate).value)

print(
    clathrate.composition.reduced_formula,
    f"{len(clathrate)} atoms, volume {clathrate.volume:.0f} A^3",
)
print(f"  true rung 3      -> mesh {rungs[3 - 1]}")
print(f"  predicted rung {predicted} -> mesh {rungs[predicted - 1]}")
clathrate = Structure.from_file(SNAPSHOT / "56721.cif")
rungs = ladder(clathrate)
predicted = int(model.predict(clathrate).value)

print(
    clathrate.composition.reduced_formula,
    f"{len(clathrate)} atoms, volume {clathrate.volume:.0f} A^3",
)
print(f"  true rung 3      -> mesh {rungs[3 - 1]}")
print(f"  predicted rung {predicted} -> mesh {rungs[predicted - 1]}")

Si 34 atoms, volume 800 A^3
  true rung 3      -> mesh (3, 3, 3)
  predicted rung 3 -> mesh (3, 3, 3)

A crystal it does not¶

Diamond silicon in its two-atom primitive cell — the most standard test case in plane-wave DFT.

In [5]:

Copied!





a = 5.431
silicon = Structure(
    Lattice([[0, a / 2, a / 2], [a / 2, 0, a / 2], [a / 2, a / 2, 0]]),
    ["Si", "Si"],
    [[0, 0, 0], [0.25, 0.25, 0.25]],
)
prediction = model.predict(silicon)
rungs = ladder(silicon)
mesh = rungs[int(prediction.value) - 1]
low, high = prediction.details["interval"]

print(
    f"predicted rung {prediction.value:.0f} -> mesh {mesh},",
    f"{mesh[0] * mesh[1] * mesh[2]} k-points",
)
print(f"the model's own interval spans rungs {low:.0f} to {high:.0f}")
print()
print("common practice for this cell is (8, 8, 8), 512 k-points")
a = 5.431
silicon = Structure(
    Lattice([[0, a / 2, a / 2], [a / 2, 0, a / 2], [a / 2, a / 2, 0]]),
    ["Si", "Si"],
    [[0, 0, 0], [0.25, 0.25, 0.25]],
)
prediction = model.predict(silicon)
rungs = ladder(silicon)
mesh = rungs[int(prediction.value) - 1]
low, high = prediction.details["interval"]

print(
    f"predicted rung {prediction.value:.0f} -> mesh {mesh},",
    f"{mesh[0] * mesh[1] * mesh[2]} k-points",
)
print(f"the model's own interval spans rungs {low:.0f} to {high:.0f}")
print()
print("common practice for this cell is (8, 8, 8), 512 k-points")

predicted rung 38 -> mesh (38, 38, 38), 54872 k-points
the model's own interval spans rungs 4 to 34

common practice for this cell is (8, 8, 8), 512 k-points

The interval covers most of the ladder, which is one way the model can say it does not know. The published value is also well above this cell's own median (13), since the model's median already falls in the band this policy serves at q0.975 rather than at the median -- a wide interval here compounds with a policy that already publishes conservatively.

Why¶

Nothing like this cell is in the training set.

In [6]:

Copied!





sizes, seen_rungs = [], []
for sample_id, rung, formula in rows:
    if formula != "Si":
        continue
    sizes.append(len(Structure.from_file(SNAPSHOT / f"{sample_id}.cif")))
    seen_rungs.append(int(rung))

print(f"silicon entries in the training record: {len(sizes)}")
print(f"  atoms per cell: {min(sizes)} to {max(sizes)}")
print(f"  rungs:          {min(seen_rungs)} to {max(seen_rungs)}")
print(f"  any with 2 atoms: {2 in sizes}")
sizes, seen_rungs = [], []
for sample_id, rung, formula in rows:
    if formula != "Si":
        continue
    sizes.append(len(Structure.from_file(SNAPSHOT / f"{sample_id}.cif")))
    seen_rungs.append(int(rung))

print(f"silicon entries in the training record: {len(sizes)}")
print(f"  atoms per cell: {min(sizes)} to {max(sizes)}")
print(f"  rungs:          {min(seen_rungs)} to {max(seen_rungs)}")
print(f"  any with 2 atoms: {2 in sizes}")

silicon entries in the training record: 17
  atoms per cell: 4 to 58
  rungs:          3 to 25
  any with 2 atoms: False

Every silicon in the record is a large-cell allotrope. A large cell has a small Brillouin zone and needs few k-points, so the model has never seen silicon that needs many. Small cells are rare across the whole record, not only for silicon.

Cell choice¶

A rung is relative to the cell. Ask the same material twice, written two ways.

In [7]:

Copied!





conventional = Structure(
    Lattice.cubic(a),
    ["Si"] * 8,
    [
        [0, 0, 0],
        [0, 0.5, 0.5],
        [0.5, 0, 0.5],
        [0.5, 0.5, 0],
        [0.25, 0.25, 0.25],
        [0.25, 0.75, 0.75],
        [0.75, 0.25, 0.75],
        [0.75, 0.75, 0.25],
    ],
)

header = "cell".ljust(28) + "rung".rjust(6) + "mesh".rjust(15) + "k-points".rjust(10)
print(header)
for label, cell in (
    ("primitive (2 atoms)", silicon),
    ("conventional (8 atoms)", conventional),
):
    value = int(model.predict(cell).value)
    rungs = ladder(cell)
    if value > len(rungs):
        # The published rung has no mesh on this cell's own ladder.
        print(
            label.ljust(28)
            + str(value).rjust(6)
            + f"  >{rungs[-1]} (off the {len(rungs)}-rung ladder)"
        )
        continue
    grid = rungs[value - 1]
    count = grid[0] * grid[1] * grid[2]
    print(
        label.ljust(28)
        + str(value).rjust(6)
        + str(grid).rjust(15)
        + str(count).rjust(10)
    )
conventional = Structure(
    Lattice.cubic(a),
    ["Si"] * 8,
    [
        [0, 0, 0],
        [0, 0.5, 0.5],
        [0.5, 0, 0.5],
        [0.5, 0.5, 0],
        [0.25, 0.25, 0.25],
        [0.25, 0.75, 0.75],
        [0.75, 0.25, 0.75],
        [0.75, 0.75, 0.25],
    ],
)

header = "cell".ljust(28) + "rung".rjust(6) + "mesh".rjust(15) + "k-points".rjust(10)
print(header)
for label, cell in (
    ("primitive (2 atoms)", silicon),
    ("conventional (8 atoms)", conventional),
):
    value = int(model.predict(cell).value)
    rungs = ladder(cell)
    if value > len(rungs):
        # The published rung has no mesh on this cell's own ladder.
        print(
            label.ljust(28)
            + str(value).rjust(6)
            + f"  >{rungs[-1]} (off the {len(rungs)}-rung ladder)"
        )
        continue
    grid = rungs[value - 1]
    count = grid[0] * grid[1] * grid[2]
    print(
        label.ljust(28)
        + str(value).rjust(6)
        + str(grid).rjust(15)
        + str(count).rjust(10)
    )

cell                          rung           mesh  k-points
primitive (2 atoms)             38   (38, 38, 38)     54872

conventional (8 atoms)          40  >(38, 38, 38) (off the 38-rung ladder)

The conventional cell is four times larger, so it needs a coarser mesh than the primitive one -- but the model gets the ordering backwards, and worse:

⚠️ The conventional cell's predicted rung is not on its own ladder. Rung 40 exceeds the 38-rung ladder this 8-atom cubic cell has under the 0.03 Å⁻¹ floor model.json declares -- there is no mesh at rung 40 for it. Always check a published rung against the structure's own ladder length before indexing into it.

QRF95, asked the same question, returns one k-distance for both cells -- 0.217 Å⁻¹ -- because a k-distance does not depend on how the cell was written. The meshes that follow, (10, 10, 10) and (6, 6, 6), get both the ordering and the magnitude right.

What to take from this¶

A rung is meaningless without its ladder, and can exceed it. Rebuild the ladder for your structure and check the published rung against its length before indexing into it.
Small, simple cells are outside this model's training distribution. The primitive cell's recommended mesh here is over 100x denser than common practice. Sanity-check small-cell recommendations rather than trusting them unchecked.
Check the interval, not just the published rung. A wide interval (as here) is the model telling you it is unsure -- treat the number as a lower bound, not an answer, and verify convergence directly.