script dump

riso_calc.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+
+import argparse
+import sys
+import numpy as np
+import gemmi
+
+
+def read_mtz_amplitudes(mtz_path, f_label):
+    """
+    Reads an MTZ file and returns a dictionary:
+        {(h,k,l): amplitude}
+    """
+    mtz = gemmi.read_mtz_file(mtz_path)
+    mtz.setup_batched()
+
+    try:
+        col = mtz.column_with_label(f_label)
+    except RuntimeError:
+        sys.exit(f"ERROR: Column '{f_label}' not found in {mtz_path}")
+
+    hkl_dict = {}
+
+    for row in mtz:
+        h, k, l = int(row[0]), int(row[1]), int(row[2])
+        fval = row[col.idx]
+
+        if not np.isnan(fval):
+            hkl_dict[(h, k, l)] = float(fval)
+
+    return hkl_dict
+
+
+def match_reflections(dict1, dict2):
+    """
+    Returns matched amplitude arrays for common HKLs.
+    """
+    common_hkls = set(dict1.keys()) & set(dict2.keys())
+
+    if len(common_hkls) == 0:
+        sys.exit("ERROR: No matching HKLs between files.")
+
+    f1 = np.array([dict1[hkl] for hkl in common_hkls])
+    f2 = np.array([dict2[hkl] for hkl in common_hkls])
+
+    return f1, f2, len(common_hkls)
+
+
+def compute_riso(f1, f2):
+    """
+    Computes Riso:
+        Riso = sum | |F1| - |F2| | / sum |F1|
+    """
+    numerator = np.sum(np.abs(np.abs(f1) - np.abs(f2)))
+    denominator = np.sum(np.abs(f1))
+
+    if denominator == 0:
+        sys.exit("ERROR: Denominator is zero.")
+
+    return numerator / denominator
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Compute Riso between two MTZ files."
+    )
+
+    parser.add_argument("mtz1", help="Reference MTZ file")
+    parser.add_argument("mtz2", help="Second MTZ file")
+    parser.add_argument("--f1", required=True,
+                        help="Amplitude column label in first MTZ")
+    parser.add_argument("--f2", required=True,
+                        help="Amplitude column label in second MTZ")
+
+    args = parser.parse_args()
+
+    print("Reading MTZ files...")
+    dict1 = read_mtz_amplitudes(args.mtz1, args.f1)
+    dict2 = read_mtz_amplitudes(args.mtz2, args.f2)
+
+    print("Matching reflections...")
+    f1, f2, nmatch = match_reflections(dict1, dict2)
+
+    print(f"Matched reflections: {nmatch}")
+
+    riso_value = compute_riso(f1, f2)
+
+    print(f"\nRiso = {riso_value:.6f}")
+
+
+if __name__ == "__main__":
+    main()