From 0c9a90f77d4b701c404fb1aac068fba323793bd2 Mon Sep 17 00:00:00 2001
From: Beale John Henry <beale_j@ra-l-006.psi.ch>
Date: Thu, 8 May 2025 00:37:38 +0200
Subject: [PATCH] debugged and added all argparse

---
 reduction_tools/stream_select.py | 42 ++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 15 deletions(-)

diff --git a/reduction_tools/stream_select.py b/reduction_tools/stream_select.py
index 4f7d46a..4f8bfbc 100644
--- a/reduction_tools/stream_select.py
+++ b/reduction_tools/stream_select.py
@@ -110,6 +110,9 @@ def extract_xtals( chunk ):
             cell_df = scrub_cells( line )
             cells_df = pd.concat( ( cells_df, cell_df ) )
 
+    # reset index of cells_df
+    cells_df = cells_df.reset_index( drop=True )
+
     return xtals, cells_df
 
 def extract_header( chunk ):
@@ -278,7 +281,7 @@ def main( input_file, output, cell_tolerance, angle_tolerance, ideal_centre ):
     print( "done" )
 
 def list_of_floats(arg):
-    return list(map(int, arg.split(',')))
+    return list(map(float, arg.split(',')))
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -292,23 +295,32 @@ if __name__ == "__main__":
     parser.add_argument(
         "-o",
         "--output",
-        help="output stream file with sampled xtals",
+        help="output stream file name. '.stream will be added'",
         type=str,
         default="sample"
     )
-    # parser.add_argument(
-    #     "-n",
-    #     "--sample",
-    #     help="size of sample to take from input.stream",
-    #     type=int,
-    #     required=True
-    # )
-
+    parser.add_argument(
+        "-c",
+        "--cell_tolerance",
+        help="tolerance for cell lengths. 0.2 is default.",
+        type=float,
+        default=0.2
+    )
+    parser.add_argument(
+        "-a",
+        "--angle_tolerance",
+        help="tolerance for cell angles. 0.25 is default.",
+        type=float,
+        default=0.25
+    )
+    parser.add_argument(
+        "-i",
+        "--ideal_centre",
+        help="ideal lengths to be selected around. List of floats - e.g. 78.5, 78.5, 38.45, 90, 90, 90",
+        type=list_of_floats,
+        required=True
+    )
     args = parser.parse_args()
 
-    cell_tolerance = 0.2
-    angle_tolerance = 0.5
-    ideal_centre = [ 78.5, 78.5, 38.45, 90, 90, 90 ]
-
     # run main
-    main( args.stream, args.output, cell_tolerance, angle_tolerance, ideal_centre )
\ No newline at end of file
+    main( args.stream, args.output, args.cell_tolerance, args.angle_tolerance, args.ideal_centre )
\ No newline at end of file