Jungfraujoch/tools/battery/shelx_check.py

#!/usr/bin/env python3
"""SHELXL refinement of a small-molecule set's reference structure (COD) against rugnux's p.hkl.

Ground truth for small molecules, the way a chemical crystallographer would judge the data: the
published structure is refined against our intensities with a fixed recipe, and R1, wR2, GooF,
extinction, the weighting scheme and the residual density are read back. Reported, never scored.

Recipe (the same for every run, so two runs compare):
  - the data are reindexed into the COD setting: every integer matrix (entries -1..1, det +1) that
    maps the data cell onto the COD cell within 3% is tried with a 6-cycle refinement, and the one
    with the lowest R1 is kept (lowest by a wide margin on every set so far);
  - non-H atoms anisotropic, H atoms fixed at the COD positions with their COD U, occupancies
    fixed at the COD values, EXTI refined;
  - weights: WGHT 0.1, then three rounds of 12 cycles each with the weights SHELXL suggests;
  - MERG 2 (SHELXL merges equivalents itself, Friedel mates too for a centrosymmetric group), so
    the same recipe takes a merged or an unmerged HKLF 4 file; R(int)/R(sigma) are read where
    SHELXL prints them (meaningful for unmerged data only).
  - "fixed-model R1(F)": the data scored against |Fc| of the COD model as published (gemmi, no
    refinement of the model, a scale and SHELXL's extinction coefficient fitted): a check that does
    not depend on what SHELXL can absorb into ADPs. Extinction is a property of this crystal, not of
    the model, and without it the strongest reflections dominate the score.

Derived from the small-molecule harness of the 2026-10 data-quality work (codref.py/fixedmodel.py).
SHELXL is called, not shipped: give its path as site key "shelxl" or have it on PATH (it comes with
CCP4). The COD entries are downloaded once into the site's "cod_cache".
"""
import itertools
import os
import re
import shutil
import subprocess
import sys
import tempfile
import urllib.request

import gemmi
import numpy as np

DEFAULT_CACHE = "/data/battery/cod_cache"
COD_URL = "https://www.crystallography.net/cod/{}.cif"

KEYS = ("shx_status", "shx_reason", "shx_cod", "shx_r1", "shx_n4", "shx_r1_all", "shx_wr2",
        "shx_goof", "shx_exti", "shx_wght_a", "shx_wght_b", "shx_peak", "shx_hole", "shx_rint",
        "shx_rsigma", "shx_k_top", "shx_fixed_r1")


def find_shelxl(site_path=None):
    if site_path and os.path.exists(site_path):
        return site_path
    return shutil.which("shelxl")


def cod_cif(cod_id, cache):
    """The COD entry's CIF, downloaded into the cache on first use."""
    os.makedirs(cache, exist_ok=True)
    path = os.path.join(cache, f"{cod_id}.cif")
    if not os.path.exists(path):
        data = urllib.request.urlopen(COD_URL.format(cod_id), timeout=60).read()
        tmp = path + ".part"
        with open(tmp, "wb") as fh:
            fh.write(data)
        os.replace(tmp, path)
    return path


def read_report(path):
    out = {}
    for line in open(path, errors="replace"):
        if "=" in line:
            k, v = line.split("=", 1)
            out[k.strip()] = v.strip()
    return out


def read_hkl(path):
    """[(h, k, l, rest of line)] of an HKLF 4 file, up to its 0 0 0 terminator."""
    out = []
    for line in open(path):
        if len(line) < 12 or not line.strip():
            continue
        h, k, l = int(line[0:4]), int(line[4:8]), int(line[8:12])
        if h == k == l == 0:
            break
        out.append((h, k, l, line[12:].rstrip("\n")))
    return out


def write_hkl(refl, M, path):
    with open(path, "w") as fh:
        for h, k, l, rest in refl:
            n = M @ np.array([h, k, l])
            fh.write("%4d%4d%4d%s\n" % (n[0], n[1], n[2], rest))
        fh.write("   0   0   0    0.00    0.00\n")


def cell_after(dcell, M):
    """Cell of the axes a' = M a (rows of M in the data's axes)."""
    A = np.array(dcell.orth.mat.tolist())
    N = A @ np.array(M).T
    length = [np.linalg.norm(N[:, i]) for i in range(3)]
    ang = lambda u, v: np.degrees(np.arccos(np.dot(u, v) / np.linalg.norm(u) / np.linalg.norm(v)))
    return length + [ang(N[:, 1], N[:, 2]), ang(N[:, 0], N[:, 2]), ang(N[:, 0], N[:, 1])]


def spacegroup(st):
    return gemmi.find_spacegroup_by_name(st.spacegroup_hm)


def reindex_candidates(st, dcell):
    ref = st.cell.parameters
    rots = [np.array(o.rot) // 24 for o in spacegroup(st).operations().sym_ops]
    out = []
    for e in itertools.product((-1, 0, 1), repeat=9):
        M = np.array(e).reshape(3, 3)
        if round(np.linalg.det(M)) != 1:
            continue
        c = cell_after(dcell, M)
        if all(abs(c[i] - ref[i]) / ref[i] < 0.03 for i in range(3)) and \
                all(abs(c[i] - ref[i]) < 2.0 for i in range(3, 6)):
            # one per class of matrices equivalent under the model's point group
            if not any(any(np.array_equal(R @ M, U) for R in rots) for U, _ in out):
                out.append((M, c))
    return out


def write_ins(st, lam, cellp, wght, cycles, path):
    sg = spacegroup(st)
    ops = sg.operations()
    centro = ops.is_centrosymmetric()
    latt = {"P": 1, "I": 2, "R": 3, "F": 4, "A": 5, "B": 6, "C": 7}[sg.hm[0]]
    symm = []
    for o in ops.sym_ops:
        R = np.array(o.rot) // 24
        if np.array_equal(R, np.eye(3, dtype=int)) and not any(o.tran):
            continue
        if centro and round(np.linalg.det(R)) < 0:
            continue
        symm.append(o.triplet().upper())
    counts = {}
    for s in st.get_all_unit_cell_sites():
        counts[s.element.name] = counts.get(s.element.name, 0) + s.occ
    elems = sorted(counts, key=lambda e: (e != "C", e != "H", e))
    lines = ["TITL cod", "CELL %.5f %.4f %.4f %.4f %.3f %.3f %.3f" % (lam, *cellp),
             "ZERR 1 0.001 0.001 0.001 0.01 0.01 0.01", "LATT %d" % (latt if centro else -latt)]
    lines += ["SYMM " + s for s in symm]
    lines += ["SFAC " + " ".join(elems), "UNIT " + " ".join(str(max(1, round(counts[e]))) for e in elems),
              f"L.S. {cycles}", "LIST 4", "FMAP 2", "PLAN 5", "MERG 2", "EXTI 0", "WGHT " + wght,
              "FVAR 1.0"]
    names, atoms, nonh = set(), [], []
    for i, s in enumerate(st.sites):
        name = re.sub(r"[^A-Za-z0-9]", "", s.label)[:4]
        if not name or name in names:            # SHELXL names: 4 characters, unique
            name = f"{s.element.name[:1]}{i:03d}"[:4]
        names.add(name)
        sf = elems.index(s.element.name) + 1
        x, y, z = s.fract.x, s.fract.y, s.fract.z
        occ = 10.0 + (s.occ if s.occ > 0 else 1.0)
        if s.element.name == "H":
            u = s.u_iso if s.u_iso > 0 else 0.05
            atoms.append("%-4s %d %9.5f %9.5f %9.5f %8.5f %8.5f" % (name, sf, 10 + x, 10 + y, 10 + z, occ, 10 + u))
        else:
            u = s.u_iso if s.u_iso > 0 else 0.02
            atoms.append("%-4s %d %9.5f %9.5f %9.5f %8.5f %8.5f" % (name, sf, x, y, z, occ, u))
            nonh.append(name)
    lines += ["ANIS " + " ".join(nonh[i:i + 12]) for i in range(0, len(nonh), 12)]
    lines += atoms + ["HKLF 4", "END"]
    with open(path, "w") as fh:
        fh.write("\n".join(lines) + "\n")


def run_shelxl(shelxl, name, cwd):
    env = dict(os.environ, OMP_NUM_THREADS="4")
    subprocess.run([shelxl, name], cwd=cwd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                   env=env, timeout=900)
    lst_p, res_p = os.path.join(cwd, name + ".lst"), os.path.join(cwd, name + ".res")
    lst = open(lst_p, errors="replace").read() if os.path.exists(lst_p) else ""
    res = open(res_p, errors="replace").read() if os.path.exists(res_p) else ""
    last = lambda pat, text=lst: (re.findall(pat, text, re.M) or [None])[-1]
    num = lambda s: float(s) if s is not None else None
    wght = (last(r"^WGHT\s+(.*)$", res) or "").split()
    k_top = None
    i = lst.rfind("Analysis of variance for reflections employed in refinement")
    if i >= 0:
        m = re.search(r"\n\s+K\s+([-\d. ]+)", lst[i:i + 3000])
        k_top = num(m.group(1).split()[-1]) if m else None
    n4 = last(r"R1 =\s*[\d.]+ for\s+(\d+) Fo > 4sig")
    return {"shx_r1": num(last(r"R1 =\s*([\d.]+) for\s+\d+ Fo > 4sig")),
            "shx_n4": int(n4) if n4 else None,
            "shx_r1_all": num(last(r"and\s+([\d.]+) for all")),
            "shx_wr2": num(last(r"wR2 =\s*([\d.]+)")),
            "shx_goof": num(last(r"GooF = S =\s*([\d.]+)")),
            "shx_peak": num(last(r"Highest peak\s+([\d.-]+)")),
            "shx_hole": num(last(r"Deepest hole\s+([\d.-]+)")),
            "shx_rint": num(last(r"R\(int\) =\s*([\d.]+)")),
            "shx_rsigma": num(last(r"R\(sigma\) =\s*([\d.]+)")),
            "shx_exti": num(last(r"^EXTI\s+([\d.]+)", res)),
            "shx_wght_a": num(wght[0]) if wght else None,
            "shx_wght_b": num(wght[1]) if len(wght) > 1 else None,
            "shx_k_top": k_top}, " ".join(wght)


def fixed_model_r1(st, refl, M, lam):
    """R1(F) of the (reindexed) data against |Fc| of the COD model as published, with one scale and
    SHELXL's extinction, Fc2* = Fc2 (1 + 0.001 x Fc2 lambda^3 / sin 2theta)^-1/2, x fitted by a scan."""
    sg = spacegroup(st)
    ops = sg.operations()
    asu = gemmi.ReciprocalAsu(sg)
    calc = gemmi.StructureFactorCalculatorX(st.cell)
    st.setup_cell_images()
    groups = {}
    for h, k, l, rest in refl:
        n = M @ np.array([h, k, l])
        if ops.is_systematically_absent(list(map(int, n))):
            continue
        p = rest.split()
        i, s = float(p[0]), float(p[1])
        if s <= 0:
            continue
        groups.setdefault(tuple(asu.to_asu(list(map(int, n)), ops)[0]), []).append((i, s))
    if not groups:
        return None
    fo, fc, ext = [], [], []
    for hkl, obs in groups.items():
        w = [1 / s ** 2 for _, s in obs]
        fo.append(sum(i * wi for (i, _), wi in zip(obs, w)) / sum(w))
        fc.append(abs(calc.calculate_sf_from_small_structure(st, list(hkl))) ** 2)
        sin_th = lam / (2 * st.cell.calculate_d(list(hkl)))
        ext.append(0.001 * lam ** 3 / (2 * sin_th * np.sqrt(1 - sin_th ** 2)))
    fo, fc, ext = np.array(fo), np.array(fc), np.array(ext)

    def r1(x):
        fcx = fc / np.sqrt(1 + x * fc * ext)
        f = fo / (np.sum(fo * fcx) / np.sum(fcx * fcx))
        Fo = np.sqrt(np.clip(f, 0, None))
        return float(np.sum(np.abs(Fo - np.sqrt(fcx))) / np.sum(Fo))
    return round(min(r1(x) for x in np.concatenate(([0.0], np.geomspace(1e-4, 100, 121)))), 4)


def check(wd, cod_id, cache=DEFAULT_CACHE, shelxl=None, work=None):
    """shx_* fields for one finished set in work dir `wd` (p.hkl + p_report.txt). `work`: where
    SHELXL runs (default <wd>/shelx)."""
    out = {k: None for k in KEYS}
    out["shx_cod"] = cod_id
    if not shelxl:
        return dict(out, shx_status="skipped", shx_reason="SHELXL not available (site key 'shelxl' or PATH)")
    hkl, rep_p = os.path.join(wd, "p.hkl"), os.path.join(wd, "p_report.txt")
    if not os.path.exists(hkl) or not os.path.exists(rep_p):
        return dict(out, shx_status="skipped", shx_reason="no p.hkl")
    try:
        rep = read_report(rep_p)
        dcell = gemmi.UnitCell(*map(float, rep["UNIT_CELL_CONSTANTS"].split()))
        lam = float(rep["WAVELENGTH"])
        st = gemmi.read_small_structure(cod_cif(cod_id, cache))
        refl = read_hkl(hkl)
        cands = reindex_candidates(st, dcell)
        if not cands:
            return dict(out, shx_status="skipped", shx_reason="data cell does not match the COD cell")
        work = work or os.path.join(wd, "shelx")
        os.makedirs(work, exist_ok=True)
        best = None
        for j, (M, c) in enumerate(cands):
            d = os.path.join(work, f"trial{j}")
            os.makedirs(d, exist_ok=True)
            write_hkl(refl, M, os.path.join(d, "t.hkl"))
            write_ins(st, lam, c, "0.1", 6, os.path.join(d, "t.ins"))
            r1 = run_shelxl(shelxl, "t", d)[0]["shx_r1"]
            if r1 is not None and (best is None or r1 < best[0]):
                best = (r1, M, c)
        if best is None:
            return dict(out, shx_status="error", shx_reason="SHELXL gave no R1 in any setting")
        _, M, c = best
        d = os.path.join(work, "final")
        os.makedirs(d, exist_ok=True)
        write_hkl(refl, M, os.path.join(d, "f.hkl"))
        wght = "0.1"
        for _ in range(3):
            write_ins(st, lam, c, wght, 12, os.path.join(d, "f.ins"))
            vals, suggested = run_shelxl(shelxl, "f", d)
            wght = suggested or wght
        out.update(vals)
        out["shx_fixed_r1"] = fixed_model_r1(st, refl, M, lam)
        out["shx_status"] = "ok" if vals["shx_r1"] is not None else "error"
        if vals["shx_r1"] is None:
            out["shx_reason"] = "SHELXL did not finish the refinement (see shelx/final/f.lst)"
    except Exception as ex:                   # a failing check must not lose the set's result
        out.update(shx_status="error", shx_reason=f"{type(ex).__name__}: {ex}"[:300])
    return out


def check_tmp(wd, cod_id, cache, shelxl):
    """check() for a finished (read-only) run directory: SHELXL works in a temporary directory."""
    with tempfile.TemporaryDirectory(prefix="shelx_") as tmp:
        return check(wd, cod_id, cache, shelxl, tmp)


if __name__ == "__main__":
    # shelx_check.py <set work dir> <COD id> [cod cache]
    print(check_tmp(sys.argv[1], sys.argv[2], sys.argv[3] if len(sys.argv) > 3 else DEFAULT_CACHE,
                    find_shelxl()))