updated file output + number of crystals to table

cool_tools/table1.py

@@ -11,18 +11,22 @@ itself and isn't available here.
 
 Usage:
   python3 table1.py data/dtpaa_20fs_high --stream_file data/dtpaa_20fs_high_76-77_0.74-scaled.stream
-  python3 table1.py data/hewl_nBW data/hewl_LBW --stream_file data/nBW.stream data/LBW.stream -o table1.tsv
+  python3 table1.py data/hewl_nBW data/hewl_LBW --stream_file data/nBW.stream data/LBW.stream -o table1.csv
 """
 
 import argparse
+import csv
 import os
 import re
 
+HEADER = "Reflections measured after indexing"
+
 ROW_ORDER = [
     "Wavelength",
     "Resolution range",
     "Space group",
     "Unit cell",
+    "Number of crystals",
     "Total reflections",
     "Unique reflections",
     "Multiplicity",
@@ -82,6 +86,11 @@ def wavelength_from_stream(stream_path):
     return None
 
 
+def count_crystals(stream_path):
+    with open(stream_path) as f:
+        return sum(1 for line in f if line.strip() == HEADER)
+
+
 def gather_stats(proc_dir, stream_file=None):
     check_hkl_log = _read(os.path.join(proc_dir, "check_hkl.log"))
     cc_log = _read(os.path.join(proc_dir, "cc.log"))
@@ -121,21 +130,23 @@ def gather_stats(proc_dir, stream_file=None):
     overall_ccstar = _search(ccstar_log, r"Overall CC\* = ([\d.]+)")
 
     wavelength = wavelength_from_stream(stream_file) if stream_file else None
+    n_crystals = count_crystals(stream_file) if stream_file else None
 
     return {
-        "Wavelength": f"{wavelength:.4f}" if wavelength else "N/A",
+        "Wavelength": f"{wavelength:.2f}" if wavelength else "N/A",
         "Resolution range": f"{low_res:.2f} - {high_res:.2f} ({shell_low:.2f} - {shell_high:.2f})",
         "Space group": spacegroup or "N/A",
-        "Unit cell": " ".join(f"{v:.3f}" for v in cell) if cell else "N/A",
+        "Unit cell": " ".join(f"{v:.2f}" for v in cell) if cell else "N/A",
+        "Number of crystals": f"{n_crystals:,}" if n_crystals is not None else "N/A",
         "Total reflections": f"{total_measurements:,}" if total_measurements is not None else "N/A",
         "Unique reflections": f"{total_unique:,} ({shell_unique:,})" if total_unique is not None else "N/A",
-        "Multiplicity": f"{overall_mult:.6f} ({shell_mult:.1f})" if overall_mult is not None else "N/A",
+        "Multiplicity": f"{overall_mult:.2f} ({shell_mult:.2f})" if overall_mult is not None else "N/A",
         "Completeness (%)": f"{overall_comp:.2f} ({shell_comp:.2f})" if overall_comp is not None else "N/A",
-        "Mean I/sigma(I)": f"{overall_snr:.6f} ({shell_snr:.2f})" if overall_snr is not None else "N/A",
+        "Mean I/sigma(I)": f"{overall_snr:.2f} ({shell_snr:.2f})" if overall_snr is not None else "N/A",
         "Wilson B-factor": f"{wilson_b:.2f}" if wilson_b is not None else "N/A",
         "R-split": f"{overall_rsplit:.2f} ({shell_rsplit:.2f})" if overall_rsplit is not None else "N/A",
-        "CC1/2": f"{overall_cc:.7f} ({shell_cc:.7f})" if overall_cc is not None else "N/A",
-        "CC*": f"{overall_ccstar:.7f} ({shell_ccstar:.7f})" if overall_ccstar is not None else "N/A",
+        "CC1/2": f"{overall_cc:.2f} ({shell_cc:.2f})" if overall_cc is not None else "N/A",
+        "CC*": f"{overall_ccstar:.2f} ({shell_ccstar:.2f})" if overall_ccstar is not None else "N/A",
     }
 
 
@@ -151,11 +162,12 @@ def main():
     parser.add_argument(
         "--stream_file", nargs="+",
         help="The .stream file for each proc_dir, in the same order, used to read the "
-        "wavelength (from photon_energy). Omit a directory's wavelength by passing "
-        "'' in its place. If not given at all, Wavelength is left as N/A.",
+        "wavelength (from photon_energy) and count the number of crystals. Omit a "
+        "directory's by passing '' in its place. If not given at all, both are "
+        "left as N/A.",
     )
     parser.add_argument(
-        "-o", "--output", help="Write the table as tab-separated values to this file."
+        "-o", "--output", help="Write the table as a csv file to this path."
     )
     args = parser.parse_args()
 
@@ -171,17 +183,15 @@ def main():
         label = os.path.basename(os.path.normpath(proc_dir))
         columns[label] = gather_stats(proc_dir, stream_file or None)
 
-    header = "\t".join([""] + list(columns.keys()))
-    lines = [header]
+    rows = [[""] + list(columns.keys())]
     for row in ROW_ORDER:
-        lines.append("\t".join([row] + [columns[label][row] for label in columns]))
-    table_text = "\n".join(lines)
+        rows.append([row] + [columns[label][row] for label in columns])
 
-    print(table_text)
+    print("\n".join(", ".join(row) for row in rows))
 
     if args.output:
-        with open(args.output, "w") as f:
-            f.write(table_text + "\n")
+        with open(args.output, "w", newline="") as f:
+            csv.writer(f).writerows(rows)
         print(f"\nWrote table to {args.output}")