"""Investigate the low performance of the ``synthmorph`` pipeline on the M6 CT->T1c rigid benchmark (see REGISTRATION_BENCHMARK_RIGID.md). The benchmark median MI is 0.277 for synthmorph vs ~0.65 for the intensity tools, with the ~30 deg repositioned cases collapsing to 0.06-0.11. This script decomposes *why*, scoring everything with the exact benchmark scoring (``scripts/bench_registration.py::evaluate``): A. Transform decomposition -- for each case, run synthmorph ``rigid`` (the benchmark config; deterministic, so it reproduces the bench JSON) and a ``simpleitk`` benchmark reference, and report per case: - achieved MI / ImpactLoss / skull-ring gap for both, - rotation angle + axis of each transform (world RAS frame), - the *relative* rotation error angle(R_ref^T R_sm), - the head-centre placement error |T_sm(c) - T_ref(c)| (c = T1c head centroid), - the identity-transform score (the "do nothing" baseline), - a numeric rigidity check of the synthmorph matrix (R^T R ~ I, det = 1). B. Affine model -- run synthmorph ``affine`` (12-DOF matrix model) on the same cases and score it: does the bigger matrix family recover the tilt that rigid does not? The linear part is decomposed into rotation / scale / shear (polar). C. Joint (affine + deformable) upper bound -- for the worst cases, run ``-m joint`` directly in Docker (outside the rigid contract: the output is a displacement field, applied with the tool's own ``-o`` resampling) and score the warped CT image. Answers: is the low score a matrix-model limitation or a model-family (pretrained-on-MRI) limitation? D. Network-space mechanics -- static (no registration) check of what the network actually sees, emulating ``network_space`` + ``normalize`` from the image code (isotropic 1-mm LIA grid, extent 256, centred on each volume's own centre, fill 0, min-max into [0,1]): - head truncation: fraction of head-mask voxels (the T1c head mapped into CT space via the reference transform) outside the 256 mm box, for the CT and T1c network spaces, - intensity band: the min-max-normalized values over the head region for the T1c, the [0,80]-clipped CT (the benchmark default) and the raw CT. Usage (repo root, conda env longitudinal-313): python scripts/investigate_synthmorph.py --workers 3 python scripts/investigate_synthmorph.py --parts D # static only python scripts/investigate_synthmorph.py --keys m6_2BBN5XA4_20250411 """ import argparse import importlib.util import json import math import os import shutil import subprocess import sys import tempfile import time from concurrent.futures import ProcessPoolExecutor, as_completed ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, ROOT) import numpy as np import SimpleITK as sitk from scipy.ndimage import center_of_mass from src.common import d, read_nii_img, head_mask_from_image from src import registration_pipelines as RP def _load_module(name, rel): p = os.path.join(ROOT, rel) spec = importlib.util.spec_from_file_location(name, p) m = importlib.util.module_from_spec(spec) spec.loader.exec_module(m) return m BR = _load_module("bench_registration", "scripts/bench_registration.py") M6 = RP._load_m6() # the worst 3 bench cases (rigid MI): ~30 deg tilt x2 + a mid-tier case JOINT_KEYS = ("m6_2BBN5XA4_20250411", "m6_24V62QQM_20240820", "m6_24V62QQM_20230223") # --------------------------------------------------------------------------- # # transform decomposition # --------------------------------------------------------------------------- def _rot_angle(R): c = (np.trace(R) - 1.0) / 2.0 return math.degrees(math.acos(float(np.clip(c, -1.0, 1.0)))) def _rot_axis(R, theta_deg): if theta_deg <= 0.5: return np.zeros(3) axis = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]]) n = np.linalg.norm(axis) return axis / n if n > 1e-9 else np.zeros(3) def _to_affine(T): """Recover (R, t) of any SimpleITK transform by point mapping (world RAS).""" o = np.array(T.TransformPoint((0.0, 0.0, 0.0)), dtype=float) R = np.empty((3, 3), dtype=float) for i in range(3): p = [0.0, 0.0, 0.0] p[i] = 1.0 R[:, i] = np.array(T.TransformPoint(tuple(p)), dtype=float) - o return R, o def _decompose_affine(T): """(R, t, theta_deg, axis, rigidity) of a SimpleITK transform in world RAS.""" R, t = _to_affine(T) theta = _rot_angle(R) rig = {"det": round(float(np.linalg.det(R)), 6), "max_devi_from_I": round(float(np.linalg.norm(R.T @ R - np.eye(3))), 6)} return R, t, theta, _rot_axis(R, theta), rig def _angle_between_rotations(A, B): """angle of A^T B in degrees (rotation B needs on top of A).""" return _rot_angle(A.T @ B) def _polar_rotation(M): """rotation part of a general 3x3 linear map (SVD polar).""" U, _s, Vt = np.linalg.svd(M) R_p = U @ Vt if np.linalg.det(R_p) < 0: U[:, -1] *= -1 R_p = U @ Vt return R_p def head_centroid_world(t1c): ha = sitk.GetArrayFromImage(t1c).astype("float32") idx = center_of_mass(head_mask_from_image(ha)) sp = np.array(t1c.GetSpacing()) o = np.array(t1c.GetOrigin()) D = np.array(t1c.GetDirection()).reshape(3, 3) return o + D @ (np.array(idx[::-1]) * sp) def score_identity(ct, t1c): """The "do nothing" baseline: CT resampled onto the T1c grid with the identity transform (i.e. no registration at all).""" ha = sitk.GetArrayFromImage(t1c).astype("float32") wc = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, sitk.Transform(), sitk.sitkLinear)) gap = M6._skull_ring_gap(wc, ha) mi = BR._mi_masked(wc, ha, head_mask_from_image(ha)) return {"mi": mi, "gap": gap} # --------------------------------------------------------------------------- # # part A/B/D worker # --------------------------------------------------------------------------- def _case_task(case, args): out = {"key": case["key"], "subject": case["subject"], "date": case["date"]} ct = read_nii_img(case["ct"]) t1c = read_nii_img(case["img"]) tv = read_nii_img(case["label"]) if not BR.ct_complete(ct): out["note"] = "ct_incomplete (skipped, same as the benchmark)" return out gpu = (args.gpu_base + abs(hash(out["key"])) % max(1, args.n_gpus)) % max(1, args.n_gpus) parts = set(args.parts) # geometry facts for nm, img in (("ct", ct), ("t1c", t1c)): sz = img.GetSize() sp = img.GetSpacing() c = np.array(img.TransformIndexToPhysicalPoint( ((sz[0] - 1) // 2, (sz[1] - 1) // 2, (sz[2] - 1) // 2))) out[nm] = {"shape_xyz": list(sz), "spacing_xyz": [round(float(s), 3) for s in sp], "fov_mm_xyz": [round(float(s0 * s1), 1) for s0, s1 in zip(sz, sp)], "center_world": [round(float(x), 1) for x in c]} c_t1c = head_centroid_world(t1c) out["t1c_head_centroid_world"] = [round(float(x), 1) for x in c_t1c] R_ref = t_ref = None if parts & {"A", "B", "D"}: # reference: simpleitk benchmark mode (deterministic at fixed threads) T_ref, meta_ref = RP.register_simpleitk(ct, t1c, threads=16, mode="benchmark") sref = BR.evaluate(ct, t1c, tv, T_ref) R_ref, t_ref, th_ref, ax_ref, _ = _decompose_affine(T_ref) out["reference"] = { "tool": meta_ref["toolkit"], "R": R_ref.tolist(), "t": t_ref.tolist(), "mi": sref["mi_ct_mr"], "impact_loss": sref["impact_loss"], "gap": sref["skull_ring_gap_mm"], "rotation_deg": round(th_ref, 2), "axis": [round(float(a), 3) for a in ax_ref], "translation": [round(float(x), 1) for x in t_ref], "head_centre_ct": [round(float(x), 1) for x in (R_ref @ c_t1c + t_ref)], } out["identity"] = score_identity(ct, t1c) if "A" in parts: t0 = time.time() T_sm, meta_sm = RP.register_synthmorph(ct, t1c, model="rigid", gpu=int(gpu), extent=256) ssm = BR.evaluate(ct, t1c, tv, T_sm) R_sm, t_sm, th_sm, ax_sm, rig = _decompose_affine(T_sm) p_ref = R_ref @ c_t1c + t_ref p_sm = R_sm @ c_t1c + t_sm out["rigid"] = { "seconds": round(time.time() - t0, 1), "auto_inverted": meta_sm["auto_inverted"], "clip_ct": meta_sm["clip_ct"], "mi": ssm["mi_ct_mr"], "impact_loss": ssm["impact_loss"], "gap": ssm["skull_ring_gap_mm"], "rotation_deg": round(th_sm, 2), "axis": [round(float(a), 3) for a in ax_sm], "translation": [round(float(x), 1) for x in t_sm], "matrix_rigidity": rig, "rot_err_vs_ref_deg": round(_angle_between_rotations(R_ref, R_sm), 2), "head_centre_err_mm": round(float(np.linalg.norm(p_sm - p_ref)), 1), } if "B" in parts: t0 = time.time() T_af, meta_af = RP.register_synthmorph(ct, t1c, model="affine", gpu=int(gpu), extent=256) saf = BR.evaluate(ct, t1c, tv, T_af) M_af, t_af = _to_affine(T_af) svals = np.linalg.svd(M_af, compute_uv=False) lin = {"scale_mean": round(float(np.mean(svals)), 4), "scale_iso_dev": round(float(np.std(svals)) / max(float(np.mean(svals)), 1e-12), 4), "rotation_deg": round(_rot_angle(_polar_rotation(M_af)), 2), "det": round(float(np.linalg.det(M_af)), 4)} row = {"seconds": round(time.time() - t0, 1), "auto_inverted": meta_af["auto_inverted"], "mi": saf["mi_ct_mr"], "impact_loss": saf["impact_loss"], "gap": saf["skull_ring_gap_mm"], "linear": lin} if R_ref is not None: row["rot_err_vs_ref_deg"] = round( _angle_between_rotations(R_ref, _polar_rotation(M_af)), 2) row["head_centre_err_mm"] = round( float(np.linalg.norm(M_af @ c_t1c + t_af - (R_ref @ c_t1c + t_ref))), 1) out["affine"] = row if "D" in parts: out["network_space"] = _network_space_analysis(ct, t1c, out, R_ref=R_ref, t_ref=t_ref) return out # --------------------------------------------------------------------------- # # part D: what the network sees # --------------------------------------------------------------------------- def _network_grid(img, extent=256): """The matrix-model network space of ``img``: isotropic 1-mm LIA grid, ``extent``^3, centred on the volume's own centre (mirrors ``network_space`` in the container's registration.py).""" sz = img.GetSize() c = np.array(img.TransformIndexToPhysicalPoint( ((sz[0] - 1) // 2, (sz[1] - 1) // 2, (sz[2] - 1) // 2))) # LIA axes in the RAS world: x_L = -x_R, y_I = -y_A, z_A = +z_R origin = (c[0] + extent / 2.0 - 0.5, c[1] + extent / 2.0 - 0.5, c[2] - extent / 2.0 + 0.5) net = sitk.Image((extent, extent, extent), sitk.sitkFloat32) net.SetSpacing((1.0, 1.0, 1.0)) net.SetOrigin(tuple(float(x) for x in origin)) net.SetDirection((-1.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0)) return net, c def _to_network(img, net): a = sitk.GetArrayFromImage( sitk.Resample(img, net, sitk.Transform(), sitk.sitkLinear, 0.0)) lo, hi = float(a.min()), float(a.max()) if hi - lo < 1e-9: return np.zeros_like(a), lo, hi return (a - lo) / (hi - lo), lo, hi def _network_space_analysis(ct, t1c, out, extent=256, R_ref=None, t_ref=None): """Extent truncation + min-max intensity band (what the model sees).""" res = {"extent": extent} ha = sitk.GetArrayFromImage(t1c).astype("float32") head_t1c = sitk.GetImageFromArray(head_mask_from_image(ha).astype("float32")) head_t1c.CopyInformation(t1c) if R_ref is not None: T_ref = sitk.AffineTransform(3) T_ref.SetMatrix(list(np.array(R_ref).ravel())) T_ref.SetTranslation(list(np.array(t_ref))) # T_ref maps T1c -> CT, so rasterising the T1c head onto the CT grid # needs the inverse as the destination -> source map head_ct = sitk.Resample(head_t1c, ct, T_ref.GetInverse(), sitk.sitkNearestNeighbor, 0.0) else: head_ct = head_t1c # truncation: head voxels outside each volume's OWN extent box for nm, img, head in (("ct", ct, head_ct), ("t1c", t1c, head_t1c)): _net, c = _network_grid(img, extent) idx = np.argwhere(sitk.GetArrayFromImage(head) > 0.5) if idx.size == 0: res[f"{nm}_head_trunc_frac"] = None continue sp = np.array(img.GetSpacing()) o = np.array(img.GetOrigin()) D = np.array(img.GetDirection()).reshape(3, 3) world = o[None, :] + (idx[:, ::-1] * sp) @ D.T # idx is (z,y,x) half = extent / 2.0 inside = (np.abs(world - c[None, :]) <= half).all(axis=1) head_c = world.mean(axis=0) res[f"{nm}_head_voxels"] = int(idx.shape[0]) res[f"{nm}_head_trunc_frac"] = round(float(1.0 - inside.mean()), 4) res[f"{nm}_head_off_volume_centre_mm"] = \ [round(float(x), 1) for x in (head_c - c)] # min-max normalized band over the head region (emulating normalize=True) bands = {} for nm, img, head in (("t1c", t1c, head_t1c), ("ct_clip", RP._clip_intensity(ct, 0.0, 80.0), head_ct), ("ct_raw", ct, head_ct)): net, _c = _network_grid(img, extent) a, lo, hi = _to_network(img, net) hnet = sitk.GetArrayFromImage( sitk.Resample(head, net, sitk.Transform(), sitk.sitkNearestNeighbor, 0.0)) > 0.5 v = a[hnet] bands[nm] = {"src_min": round(lo, 1), "src_max": round(hi, 1), "head_nvox_net": int(v.size), "p5": round(float(np.percentile(v, 5)), 4) if v.size else None, "p50": round(float(np.percentile(v, 50)), 4) if v.size else None, "p95": round(float(np.percentile(v, 95)), 4) if v.size else None, "band_width_p95_p5": (round(float(np.percentile(v, 95) - np.percentile(v, 5)), 4) if v.size else None)} res["normalized_band"] = bands return res # --------------------------------------------------------------------------- # # part C: joint (affine + deformable) upper bound # --------------------------------------------------------------------------- def _joint_task(case, args): out = {"key": case["key"], "subject": case["subject"], "date": case["date"], "model": "joint (affine + deformable)"} ct = read_nii_img(case["ct"]) t1c = read_nii_img(case["img"]) if not BR.ct_complete(ct): out["note"] = "ct_incomplete (skipped)" return out gpu = (args.gpu_base + abs(hash(out["key"])) % max(1, args.n_gpus)) % max(1, args.n_gpus) tmp = tempfile.mkdtemp(prefix="sm_joint_") try: ct_clip = RP._clip_intensity(ct, 0.0, 80.0) sitk.WriteImage(sitk.Cast(ct_clip, sitk.sitkFloat32), os.path.join(tmp, "moving.nii.gz"), True) sitk.WriteImage(sitk.Cast(t1c, sitk.sitkFloat32), os.path.join(tmp, "fixed.nii.gz"), True) # CLI order is (moving, fixed); -o saves moving registered to fixed cmd = ["docker", "run", "--rm", "--user", f"{os.getuid()}:{os.getgid()}", "-e", "TF_FORCE_GPU_ALLOW_GROWTH=true", "--gpus", "all", "-e", f"CUDA_VISIBLE_DEVICES={gpu}", "-v", f"{tmp}:/data", RP.SYNTHMORPH_IMAGE, "register", "-m", "joint", "-e", "256", "-g", "-o", "/data/warped.nii.gz", "/data/moving.nii.gz", "/data/fixed.nii.gz"] t0 = time.time() proc = subprocess.run(cmd, capture_output=True, text=True, timeout=7200) dt = time.time() - t0 if proc.returncode != 0: out["error"] = (f"docker failed rc={proc.returncode}:\n" + "\n".join((proc.stdout + proc.stderr).splitlines()[-15:])) return out warped = sitk.GetArrayFromImage( sitk.ReadImage(os.path.join(tmp, "warped.nii.gz"))).astype("float32") ha = sitk.GetArrayFromImage(t1c).astype("float32") gap = M6._skull_ring_gap(warped, ha) mi = BR._mi_masked(warped, ha, head_mask_from_image(ha)) il = BR._impact_loss_value(warped, ha) out.update({"seconds": round(dt, 1), "mi": mi, "impact_loss": il, "gap": gap, "gpu": int(gpu)}) finally: if not args.keep_scratch: shutil.rmtree(tmp, ignore_errors=True) return out # --------------------------------------------------------------------------- # # part F: re-centred CT probe (volume-centre anchoring hypothesis) # --------------------------------------------------------------------------- def _recenter_task(case, args, ref_map): """H5 probe: the matrix model's network space is centred on each volume's OWN centre. The T1c is a tight head crop (centre ~ head centre), but the head+neck CT's volume centre sits 15-80 mm below the head. If the model under-corrects that offset, cropping the CT to the head box (re-centring the volume on the head) should remove the translation miss. The crop is world-preserving, so the returned transform scores against the FULL CT exactly like the benchmark.""" out = {"key": case["key"], "subject": case["subject"], "date": case["date"]} ct = read_nii_img(case["ct"]) t1c = read_nii_img(case["img"]) tv = read_nii_img(case["label"]) if not BR.ct_complete(ct): out["note"] = "ct_incomplete (skipped)" return out ref = ref_map.get(case["key"]) if not ref: out["error"] = "no reference transform in partABD.json" return out R = np.array(ref["R"]) t = np.array(ref["t"]) T_ref = sitk.AffineTransform(3) T_ref.SetMatrix(list(R.ravel())) T_ref.SetTranslation(list(t)) gpu = (args.gpu_base + abs(hash(out["key"])) % max(1, args.n_gpus)) % max(1, args.n_gpus) # head in CT space (T_ref maps T1c -> CT, so rasterise with its inverse) ha = sitk.GetArrayFromImage(t1c).astype("float32") head_t1c = sitk.GetImageFromArray(head_mask_from_image(ha).astype("float32")) head_t1c.CopyInformation(t1c) head_ct = sitk.Resample(head_t1c, ct, T_ref.GetInverse(), sitk.sitkNearestNeighbor, 0.0) m = sitk.GetArrayFromImage(head_ct) > 0.5 idx = np.argwhere(m) # (z, y, x) lo, hi = idx.min(0), idx.max(0) cen = (lo + hi + 1) / 2.0 half = (hi - lo) / 2.0 + 15.0 / np.array(ct.GetSpacing())[::-1] # (z,y,x) sz = np.array(ct.GetSize())[::-1] # GetSize is (x,y,z); work in (z,y,x) b0 = np.clip(np.floor(cen - half).astype(int), 0, sz - 1) b1 = np.clip(np.ceil(cen + half).astype(int), 0, sz) def _volcentre(img): s = img.GetSize() return np.array(img.TransformIndexToPhysicalPoint( ((s[0] - 1) // 2, (s[1] - 1) // 2, (s[2] - 1) // 2))) hc = T_ref.TransformPoint(tuple(float(x) for x in head_centroid_world(t1c))) out["before"] = {"head_off_volume_centre_mm": [ round(float(x), 1) for x in (np.array(hc) - _volcentre(ct))]} # SimpleITK image slicing is in (x, y, z) order; b0/b1 are (z, y, x) ct_crop = ct[int(b0[2]):int(b1[2]), int(b0[1]):int(b1[1]), int(b0[0]):int(b1[0])] out["crop"] = {"shape": list(ct_crop.GetSize()), "head_off_volume_centre_mm": [ round(float(x), 1) for x in (np.array(hc) - _volcentre(ct_crop))]} t0 = time.time() T_sm, meta_sm = RP.register_synthmorph(ct_crop, t1c, model="rigid", gpu=int(gpu), extent=256) ssm = BR.evaluate(ct, t1c, tv, T_sm) # full CT, as in the benchmark R_sm, t_sm = _to_affine(T_sm) c_t1c = head_centroid_world(t1c) out["rigid_recentered"] = { "seconds": round(time.time() - t0, 1), "auto_inverted": meta_sm["auto_inverted"], "mi": ssm["mi_ct_mr"], "impact_loss": ssm["impact_loss"], "gap": ssm["skull_ring_gap_mm"], "rot_err_vs_ref_deg": round(_angle_between_rotations(R, R_sm), 2), "head_centre_err_mm": round( float(np.linalg.norm((R_sm @ c_t1c + t_sm) - (R @ c_t1c + t))), 1), } return out # --------------------------------------------------------------------------- # # main # --------------------------------------------------------------------------- def parse_args(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--parts", default="ABCD", help="which parts to run (letters A/B/C/D/F; F = re-centred " "CT probe, needs the partABD.json reference)") ap.add_argument("--n-cases", type=int, default=10) ap.add_argument("--keys", default=None, help="comma-separated case keys") ap.add_argument("--workers", type=int, default=3) ap.add_argument("--gpu-base", type=int, default=0) ap.add_argument("--n-gpus", type=int, default=3) ap.add_argument("--out-dir", default=os.path.join(ROOT, "results", "synthmorph_investigation")) ap.add_argument("--keep-scratch", action="store_true") return ap.parse_args() def main(): args = parse_args() os.makedirs(args.out_dir, exist_ok=True) labeled, _ = M6.discover() pool = list(labeled) if args.keys: bykey = {r["key"]: r for r in pool} missing = [k for k in args.keys.split(",") if k not in bykey] if missing: raise SystemExit(f"keys not found: {missing}") cases = [bykey[k.strip()] for k in args.keys.split(",")] else: cases = pool[:args.n_cases] print(f"[investigate] {len(cases)} cases, parts={args.parts}") rows, joint_rows, recenter_rows = [], [], [] do_cases = bool(set(args.parts) & {"A", "B", "D"}) do_joint = "C" in args.parts do_recenter = "F" in args.parts ref_map = {} if do_recenter: fp0 = os.path.join(args.out_dir, "partABD.json") if os.path.exists(fp0): with open(fp0) as f: ref_map = {r["key"]: r["reference"] for r in json.load(f) if "reference" in r} with ProcessPoolExecutor(max_workers=args.workers) as ex: futs = {} if do_cases: for c in cases: futs[ex.submit(_case_task, c, args)] = ("case", c) if do_joint: for c in cases: if c["key"] in JOINT_KEYS: futs[ex.submit(_joint_task, c, args)] = ("joint", c) if do_recenter: for c in cases: futs[ex.submit(_recenter_task, c, args, ref_map)] = ("recenter", c) for fu in as_completed(futs): kind, c = futs[fu] try: r = fu.result() if kind == "joint": joint_rows.append(r) print(f" [joint] {c['key']}: mi={r.get('mi')} " f"il={r.get('impact_loss')} {r.get('seconds')}s " f"err={(r.get('error') or '')[:100]}", flush=True) elif kind == "recenter": recenter_rows.append(r) rr = r.get("rigid_recentered", {}) print(f" [recenter] {c['key']}: mi={rr.get('mi')} " f"ctr_err={rr.get('head_centre_err_mm')} " f"rot_err={rr.get('rot_err_vs_ref_deg')} " f"off_before={r.get('before', {}).get('head_off_volume_centre_mm')} " f"off_after={r.get('crop', {}).get('head_off_volume_centre_mm')} " f"note={(r.get('note') or r.get('error') or '')[:60]}", flush=True) else: rows.append(r) print(f" [case] {c['key']}: ref.mi={r.get('reference', {}).get('mi')} " f"rigid.mi={r.get('rigid', {}).get('mi')} " f"affine.mi={r.get('affine', {}).get('mi')} " f"note={r.get('note', r.get('error', ''))[:60]}", flush=True) except Exception as e: # noqa print(f" {c['key']}: ERROR {e!r}", flush=True) (joint_rows if kind == "joint" else recenter_rows if kind == "recenter" else rows).append( {"key": c["key"], "error": repr(e)}) rows.sort(key=lambda x: x["key"]) joint_rows.sort(key=lambda x: x["key"]) recenter_rows.sort(key=lambda x: x["key"]) if do_cases: fp = os.path.join(args.out_dir, "partABD.json") merged = {r["key"]: r for r in rows} if os.path.exists(fp): with open(fp) as f: for r in json.load(f): # merge missing sections (a partial run must not wipe # data from a fuller earlier run) if r["key"] in merged: for k, v in r.items(): merged[r["key"]].setdefault(k, v) else: merged[r["key"]] = r with open(fp, "w") as f: json.dump(sorted(merged.values(), key=lambda x: x["key"]), f, indent=1) print("wrote", fp) if do_joint: fp = os.path.join(args.out_dir, "partC_joint.json") with open(fp, "w") as f: json.dump(joint_rows, f, indent=1) print("wrote", fp) if do_recenter: fp = os.path.join(args.out_dir, "partF_recenter.json") with open(fp, "w") as f: json.dump(recenter_rows, f, indent=1) print("wrote", fp) if __name__ == "__main__": main()