Files
extrapolation/riso_calc.py
John Beale 15ea8f8cd5 script dump
2026-02-17 08:52:57 +01:00

93 lines
2.2 KiB
Python

#!/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()