Clean SD

plotter/detector.py

@@ -1,76 +1,76 @@
 import sys
+import os
 import numpy as np
-import matplotlib.pyplot as pl
+import matplotlib.pyplot as plt
 import matplotlib.patches as patches
 import matplotlib.cm as colormap
 import scipy.stats as st
+import qrcode
 
-try:
-    shotnumber = sys.argv[1];
-except:
-    print 'Shot number is not given as proper input parameter. Default value (17178) is used.'
-    shotnumber = '17178'
 
-try:
-    time = sys.argv[2];
-except:
-    print 'Time is not given as proper input parameter. Default value (1005) is used.'
-    time = '1005'
+def detector(shotnumber, time, runnumber, title=None, is_debug=False):
+    results_folder = os.path.join('results', f'{shotnumber}_{time}', runnumber)
 
-try:
-    runnumber = sys.argv[3];
-except:
-    print 'Runnumber is not given as proper input parameter. Default value (402) is used.'
-    runnumber = '402'
-
-results_folder = 'results/'+shotnumber+'_'+time+'/'+runnumber+'/'
-
-try:
-    C = np.loadtxt(results_folder+'detector/cellcounter.dat')
-    X = np.loadtxt(results_folder+'detector/detx')
-    Y = np.loadtxt(results_folder+'detector/dety')
-except:
-    print 'Invalid input folder: '+results_folder
+    if title is None:
+        r'COMPASS #' + shotnumber + '\n ($t = $' + time + ' ms)'
 
+    try:
+        cell_counts = np.loadtxt(os.path.join(results_folder, 'detector', 'cellcounter.dat'))
+        det_x = np.loadtxt(os.path.join(results_folder, 'detector', 'detx'))
+        det_y = np.loadtxt(os.path.join(results_folder, 'detector', 'dety'))
+    except FileNotFoundError:
+        print(f'Invalid input folder: {results_folder}')
+        return
 
-xmin = X[0,0]
-xmax = X[-1,-1]
-ymin = Y[0,0]
-ymax = Y[-1,-1]
-xmean = (X[:,0]+X[:,1])/2
-ymean = (Y[:,0]+Y[:,1])/2
-cmax = C.max()
+    xmin, xmax = det_x[0, 0], det_x[-1, -1]
+    ymin, ymax = det_y[0, 0], det_y[-1, -1]
+    xmean = (det_x[:, 0] + det_x[:, 1]) / 2
+    ymean = (det_y[:, 0] + det_y[:, 1]) / 2
+    cmax = cell_counts.max() if cell_counts.max() > 0 else 1
 
-try:
-    fig = pl.figure()
-    ax = fig.gca()
-    ax.set_xlim(xmin, xmax)
-    ax.set_ylim(ymin, ymax)
-    ax.set_aspect('equal')
+    try:
+        plt.rc('font', family='serif')
+        plt.rc('text', usetex=True)
+        fig, ax = plt.subplots()
+        ax.set_xlim(xmin, xmax)
+        ax.set_ylim(ymin, ymax)
+        ax.set_aspect('equal')
 
-    # Contourf plot
-    pl.contourf(xmean, ymean, C, cmap='afmhot', corner_mask=False)
+        # Contourf plot
+        plt.contourf(xmean, ymean, cell_counts, cmap='afmhot', corner_mask=False)
 
-    # Detector cell matrix plot
-    for i in range (X[:,0].size):
-        for j in range (Y[:,0].size):
-            rect = patches.Rectangle((X[i,0], Y[j,0]), X[i,1]-X[i,0], Y[j,1]-Y[j,0], linewidth=1,
-                                     edgecolor='k', facecolor=colormap.afmhot(C[j,i]/cmax))
-            ax.add_patch(rect)
+        # Detector cell matrix plot
+        for i in range(det_x[:, 0].size):
+            for j in range(det_y[:, 0].size):
+                rect = patches.Rectangle((det_x[i, 0], det_y[j, 0]),
+                                         det_x[i, 1] - det_x[i, 0],
+                                         det_y[j, 1] - det_y[j, 0],
+                                         linewidth=1,
+                                         edgecolor='k',
+                                         facecolor=colormap.afmhot(cell_counts[j, i] / cmax))
+                ax.add_patch(rect)
 
-    print(C)
+        if is_debug:
+            print(cell_counts)
 
-    pl.xlabel(r"$x \mathrm{ [mm]}$")
-    pl.ylabel(r"$y \mathrm{ [mm]}$")
-    pl.title(r"COMPASS #"+shotnumber+" (t="+time+" s)")
+        plt.xlabel(r"toroidal detector position [mm]")
+        plt.ylabel(r"vertical detector position [mm]")
+        plt.title(title)
 
-    try:
-        print 'Save plot to '+results_folder+'detpy2_'+shotnumber+'_'+time+'.pdf'
-        pl.savefig(results_folder+'detpy2_'+shotnumber+'_'+time+'.pdf')
-        pl.clf()
-    except:
-        print 'Unable to save to : '+results_folder
-        pl.show()
-except:
-    print 'Unable to plot: '+results_folder
+        qr_content = f'raw.githubusercontent.com/taiga-project/refs/run/{runnumber}'
+        print(f'Add QR code: {qr_content}')
+        qr_image = qrcode.make(qr_content)
+        ax_qr = fig.add_axes([0.175, 0.495, 0.08, 0.5], anchor='NE', zorder=10)
+        ax_qr.imshow(qr_image, cmap='gray')
+        ax_qr.axis('off')
 
+        try:
+            filename = f"detector_{runnumber}.pdf"
+            print(f'Save plot to {os.path.join(results_folder, filename)}')
+            plt.savefig(os.path.join(results_folder, filename), dpi=300, bbox_inches='tight')
+            plt.clf()
+        except Exception as e:
+            print(f'Unable to save to: {results_folder} ({e})')
+            plt.show()
+    except Exception as e:
+        print(f'Unable to plot: {results_folder} ({e})')