From 3cfe47e761108363e004008a4502a1828e03288e Mon Sep 17 00:00:00 2001
From: Beale John Henry <beale_j@saresc-cons-04.psi.ch>
Date: Wed, 12 Jul 2023 16:10:36 +0200
Subject: [PATCH] updated to python3 and included now an arg.parse

---
 cool_tools/plot_orientations_orthographic.py | 281 ++++++++++---------
 1 file changed, 149 insertions(+), 132 deletions(-)

diff --git a/cool_tools/plot_orientations_orthographic.py b/cool_tools/plot_orientations_orthographic.py
index b8231e5..3f1113c 100644
--- a/cool_tools/plot_orientations_orthographic.py
+++ b/cool_tools/plot_orientations_orthographic.py
@@ -1,9 +1,25 @@
-# IMPORTant stuff....
-import string
+#!/usr/bin/env python3
+
+# written by Thomas Barends
+# v. minor edits by J. Beale + conversion to python3
+
+"""
+# aim
+visualise preferential orientations of crystals processed by CrystFEL
+
+# usage
+python update-geom-from-lab6.py -f <path-to-stream-file>, -r <reindex-or-not>, -s step-size
+
+# output
+creates 3 plots showing the orientations of the crystals along the 3 axes
+"""
+
+# modules
 import numpy as np
 import matplotlib.pyplot as plt
-from mpl_toolkits.mplot3d import Axes3D  
-
+from mpl_toolkits.mplot3d import Axes3D
+import argparse
+import os
 
 def ProjectionPlot(X,Y,Z,C,title):
     f=plt.figure()
@@ -35,157 +51,158 @@ def ProjectionPlot(X,Y,Z,C,title):
     ax4.set_ylabel('Y (along jet)')
     ax4.set_aspect('equal', 'box')
 
-    f.canvas.set_window_title(title)
+    f.canvas.manager.set_window_title(title)
     f.set_size_inches(9,7)
 
     plt.tight_layout()
     plt.savefig(title+'.png')
     plt.show()
-    
-    
 
-# declare lists to read reciprocal vectors
-astar=[]
-bstar=[]
-cstar=[]
+def main( stream_file, reindex=False, step=25 ):
+    # declare lists to read reciprocal vectors
+    astar=[]
+    bstar=[]
+    cstar=[]
 
-# set path and filenames
-#pathname='/lfs/l1/agorel/2018-03-20-shock-wave-8ns-cxim1816/2-indexing-pump-probe-zaef-nomulti/'
-#filename='indexed_65T_pumpprobe_SNR4_THR100.stream.rogue.stream'
+    # open stream file, read all lines and close again
+    infile=open(stream_file,'r')
+    inlines=infile.readlines()
+    infile.close()
 
-#pathname='/lfs/l1/agorel/2018-03-20-shock-wave-8ns-cxim1816/2-indexing-probe-zaef-nomulti/'
-#filename='indexed_65T_probe_SNR4_THR100.stream'
-
-pathname='./'
-filename='FAP-dark.stream'
-
-#pathname='./'
-#filename='smalla.stream'
-
-#reindex=True
-reindex=False
-
-# open stream file, read all lines and close again
-infile=open(pathname+filename,'r')
-inlines=infile.readlines()
-infile.close()
-
-
-
-# loop over lines and extract reciprocal cell axes 
-for line in inlines:
-    if 'astar' in line:
-        elements=line.split()
-        vector=[float(elements[2]),float(elements[3]),float(elements[4])]
-        astar.append(vector)
-    
-    if 'bstar' in line:
-        elements=line.split()
-        vector=[float(elements[2]),float(elements[3]),float(elements[4])]
-        bstar.append(vector)
-
-    if 'cstar' in line:
-        elements=line.split()
-        vector=[float(elements[2]),float(elements[3]),float(elements[4])]
-        cstar.append(vector)
-
-# turn all into numpy arrays
-astararray=np.asarray(astar)
-bstararray=np.asarray(bstar)
-cstararray=np.asarray(cstar)
-
-# get number of cells from shape of array
-s=np.shape(astararray)
-numcells=s[0]
-
-# declare lists for real-space cell vector directions (i.e. normalized vectors)
-adir=[]
-bdir=[]
-cdir=[]
-
-# declare lists for vector lengths
-alen=[]
-blen=[]
-clen=[]
-
-# loop over all cells
-for n in range(numcells):
-    astar=astararray[n,:]
-    bstar=bstararray[n,:]
-    cstar=cstararray[n,:]
-    Vstar=np.dot( astar , np.cross( bstar , cstar )  )    # reciprocal cell volume
-    a=(np.cross(bstar,cstar) / Vstar )                    # calculate real-space a
-    b=(np.cross(cstar,astar) / Vstar )                    # calculate real-space b
-    c=(np.cross(astar,bstar) / Vstar )                    # calculate real-space c
-
-    al=np.linalg.norm(a)                                  # calculate lengths 
-    bl=np.linalg.norm(b)
-    cl=np.linalg.norm(c)
-
-    if al<bl and reindex==True:
-        newa=-1.0*b
-        newb=a
-        newc=c
-         
-        newal=bl
-        newbl=al
-        newcl=cl
+    # loop over lines and extract reciprocal cell axes 
+    for line in inlines:
+        if 'astar' in line:
+            elements=line.split()
+            vector=[float(elements[2]),float(elements[3]),float(elements[4])]
+            astar.append(vector)
         
-        a=newa
-        b=newb
-        c=newc
-        al=newal
-        bl=newbl
-        cl=newcl 
+        if 'bstar' in line:
+            elements=line.split()
+            vector=[float(elements[2]),float(elements[3]),float(elements[4])]
+            bstar.append(vector)
 
-    adir.append(a/al)                                     # put normalized vectors in lists
-    bdir.append(b/bl)
-    cdir.append(c/cl)
-    
-    alen.append(al)                                       # put lengths in lists
-    blen.append(bl)
-    clen.append(cl)
-    
+        if 'cstar' in line:
+            elements=line.split()
+            vector=[float(elements[2]),float(elements[3]),float(elements[4])]
+            cstar.append(vector)
 
-adir=np.asarray(adir)                                     # turn lists into numpy arrays
-bdir=np.asarray(bdir)
-cdir=np.asarray(cdir)
+    # turn all into numpy arrays
+    astararray=np.asarray(astar)
+    bstararray=np.asarray(bstar)
+    cstararray=np.asarray(cstar)
 
-alen=np.asarray(alen)
-blen=np.asarray(blen)
-clen=np.asarray(clen)
+    # get number of cells from shape of array
+    s=np.shape(astararray)
+    numcells=s[0]
 
-f,((ax1,ax2,ax3))=plt.subplots(1,3)
+    # declare lists for real-space cell vector directions (i.e. normalized vectors)
+    adir=[]
+    bdir=[]
+    cdir=[]
 
-ax1.hist(10.*alen,int(5*np.sqrt(len(alen))),normed=1,facecolor='green',edgecolor='green')
-atitle=r'$\mathit{a}$'
-ax1.set_title(atitle)
+    # declare lists for vector lengths
+    alen=[]
+    blen=[]
+    clen=[]
 
-ax2.hist(10.*blen,int(5*np.sqrt(len(blen))),normed=1,facecolor='green',edgecolor='green')
-btitle=r'$\mathit{b}$'
-ax2.set_title(btitle)
+    # loop over all cells
+    for n in range(numcells):
+        astar=astararray[n,:]
+        bstar=bstararray[n,:]
+        cstar=cstararray[n,:]
+        Vstar=np.dot( astar , np.cross( bstar , cstar )  )    # reciprocal cell volume
+        a=(np.cross(bstar,cstar) / Vstar )                    # calculate real-space a
+        b=(np.cross(cstar,astar) / Vstar )                    # calculate real-space b
+        c=(np.cross(astar,bstar) / Vstar )                    # calculate real-space c
 
-ax3.hist(10.*clen,int(5*np.sqrt(len(clen))),normed=1,facecolor='green',edgecolor='green')
-ctitle=r'$\mathit{c}$'
-ax3.set_title(ctitle)
+        al=np.linalg.norm(a)                                  # calculate lengths 
+        bl=np.linalg.norm(b)
+        cl=np.linalg.norm(c)
 
+        if al<bl and reindex==True:
+            newa=-1.0*b
+            newb=a
+            newc=c
+            
+            newal=bl
+            newbl=al
+            newcl=cl
+            
+            a=newa
+            b=newb
+            c=newc
+            al=newal
+            bl=newbl
+            cl=newcl 
 
+        adir.append(a/al)                                     # put normalized vectors in lists
+        bdir.append(b/bl)
+        cdir.append(c/cl)
+        
+        alen.append(al)                                       # put lengths in lists
+        blen.append(bl)
+        clen.append(cl)
+        
 
-f.canvas.set_window_title('Unit Cell Parameter Histograms')
-f.set_size_inches(10,5)
-plt.tight_layout()
-plt.show()
+    adir=np.asarray(adir)                                     # turn lists into numpy arrays
+    bdir=np.asarray(bdir)
+    cdir=np.asarray(cdir)
 
-step=50                                                    # plot every nth point only
+    alen=np.asarray(alen)
+    blen=np.asarray(blen)
+    clen=np.asarray(clen)
 
-# plot a, coloured by length of a (normalized to maximum length observed)
-print 'I will plot every',step,'th point,',len(adir[::step,0]),'points in total...'
+    f,((ax1,ax2,ax3))=plt.subplots(1,3)
 
-ProjectionPlot(adir[::step,0],adir[::step,1],adir[::step,2],(alen[::step]/np.max(alen)),'a_axis_direction')
+    ax1.hist(10.*alen,int(5*np.sqrt(len(alen))),density=True,facecolor='green',edgecolor='green')
+    atitle=r'$\mathit{a}$'
+    ax1.set_title(atitle)
 
-ProjectionPlot(bdir[::step,0],bdir[::step,1],bdir[::step,2],(blen[::step]/np.max(blen)),'b_axis_direction')
+    ax2.hist(10.*blen,int(5*np.sqrt(len(blen))),density=True,facecolor='green',edgecolor='green')
+    btitle=r'$\mathit{b}$'
+    ax2.set_title(btitle)
 
-ProjectionPlot(cdir[::step,0],cdir[::step,1],cdir[::step,2],(clen[::step]/np.max(clen)),'c_axis_direction')
+    ax3.hist(10.*clen,int(5*np.sqrt(len(clen))),density=True,facecolor='green',edgecolor='green')
+    ctitle=r'$\mathit{c}$'
+    ax3.set_title(ctitle)
 
+    f.canvas.manager.set_window_title('Unit Cell Parameter Histograms')
+    f.set_size_inches(10,5)
+    plt.tight_layout()
+    plt.show()
+
+    # plot a, coloured by length of a (normalized to maximum length observed)
+    print('I will plot every',step,'th point,',len(adir[::step,0]),'points in total...')
+
+    ProjectionPlot(adir[::step,0],adir[::step,1],adir[::step,2],(alen[::step]/np.max(alen)),'a_axis_direction')
+
+    ProjectionPlot(bdir[::step,0],bdir[::step,1],bdir[::step,2],(blen[::step]/np.max(blen)),'b_axis_direction')
+
+    ProjectionPlot(cdir[::step,0],cdir[::step,1],cdir[::step,2],(clen[::step]/np.max(clen)),'c_axis_direction')
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-f",
+        "--stream_file",
+        help="stream file of interest for assessment of crystal orientations",
+        type=os.path.abspath
+    )
+    parser.add_argument(
+        "-r",
+        "--reindex",
+        help="do you need to adjust the cell parameters because CrystFEL has set a > b",
+        type=bool,
+        default=False
+    )
+    parser.add_argument(
+        "-s",
+        "--step",
+        help="how many points should i plot - every 1/10, 1/25?",
+        type=int,
+        default=25
+    )
+    args = parser.parse_args()
+    # run main
+    main( args.stream_file, args.reindex, args.step )
 
-# FIN....
\ No newline at end of file