compute_mean_spin_smash.py

#!/usr/bin/env python3

# version 0.2.0 - 24/11/2021

#this version works with the files generated by compute_vorticity 0.1.2 and the hadron datafile from get_kfo.py

import fileinput
import math
import numpy as np
import sys
import os
import os.path
import pickle
import gzip
from itertools import islice
from datetime import datetime

#it decides if we use the chemical (True) or kinetic freezeout data (False)
use_chem_fo_data=True

#limits for the transverse momentum
pt_min=0.1
pt_max=3.

#limits for |rapidity|
rap_lim=1.0

#we set the parameter hbarc
hbarc=0.197326

#we set how many lines of file containing the infos about the hadrons to read at once
nlines=2000000

#we get the name of input and output files
N_input_files=len(sys.argv)-1

if(N_input_files!=3):
   print ('Syntax: python3 compute_mean_spin_smash.py <vorticity pickled file> <hadron_data_inputfile> <outputfile>')
   sys.exit(1)

vorfile=sys.argv[1]
hadfile=sys.argv[2]
outfilename=sys.argv[3]

datas=np.zeros((nlines,9),dtype=np.float64) #one multidimensional np array for hadrons with fields: t,x,y,z,pt,rapidity,Sx,Sy,Sz,Py

if(vorfile[-3:]==".gz"):
    print("Opening gzipped file "+vorfile)
    vf=gzip.open(vorfile,"rb")
else:
    print("Opening file "+vorfile)
    vf=open(vorfile,"rb")
     
data=pickle.load(vf)
vf.close()

intt,inxx,inyy,inzz,invx,invy,invz,temp,omega_tx,omega_ty,omega_tz,omega_yz,omega_zx,omega_xy=data[:]

dt=intt[1]-intt[0]
dx=inxx[1]-inxx[0]
dy=inyy[1]-inyy[0]
dz=inzz[1]-inzz[0]

nt=len(intt)
nx=len(inxx)
ny=len(inyy)
nz=len(inzz)

tmin=intt[0]-dt/2.
tmax=intt[-1]+dt/2.
xmin=inxx[0]-dx/2.
xmax=inxx[-1]+dx/2.
ymin=inyy[0]-dy/2.
ymax=inyy[-1]+dy/2.
zmin=inzz[0]-dz/2.
zmax=inzz[-1]+dz/2.

sp="    " # spaces in the output file

try:
  if(hadfile[-3:]==".gz"):
    print("Opening gzipped file "+hadfile)
    hf=gzip.open(hadfile,"r")
  else:
    print("Opening file "+hadfile)
    hf=open(hadfile,"r")
except:
    print("Sorry, but I can't open "+hadfile+", therefore I quit.")
    sys.exit(2)


#opening output file
outfile=open(outfilename,"w")

count_reads=0
count_lines=0
out_of_time=0
out_of_rap=0
out_of_pt=0
out_of_space=0
bad_values=0
index=0

pdg_id = hf.readline().split()[4]
total_events = hf.readline().split()[2]
hf.readline()
hf.readline()
hf.readline()
while(True):
    it=islice(hf,count_lines,count_lines+nlines)
    for line in it:
        count_reads=count_reads+1
        discarded=False
        indata=np.float64(line.split())
        if use_chem_fo_data:
            t, x, y, z, Ep, px, py, pz = indata[0:8]
        else:
            t, x, y, z, Ep, px, py, pz = indata[8:16]
        try:
            m = math.sqrt(Ep**2-px**2-py**2-pz**2)
        except:
            bad_values+=1
            continue

        rapidity=0.5*math.log((Ep+pz)/(Ep-pz))
        pt=math.sqrt(px**2+py**2)
        if(t>intt[-1]):
           out_of_time=out_of_time+1
           discarded=True
        if(abs(rapidity)>rap_lim):
           out_of_rap=out_of_rap+1
           discarded=True
        if((z<zmin) or (z>zmax) or (x<xmin) or (x>xmax) or (y<ymin) or (y>ymax)):
           out_of_space=out_of_space+1
           discarded=True
        if((pt<pt_min) or (pt>pt_max)):
           out_of_pt=out_of_pt+1
           discarded=True
        if discarded:
           continue

        h=int(math.floor((t-tmin)/dt))
        i=int(math.floor((x-xmin)/dx))
        j=int(math.floor((y-ymin)/dy))
        k=int(math.floor((z-zmin)/dz))
        if(math.isfinite(omega_tx[h,i,j,k])):
           otx=omega_tx[h,i,j,k]
        else:
           continue
        if(math.isfinite(omega_ty[h,i,j,k])):
           oty=omega_ty[h,i,j,k]
        else:
           continue
        if(math.isfinite(omega_tz[h,i,j,k])):
           otz=omega_tz[h,i,j,k]
        else:
           continue
        if(math.isfinite(omega_yz[h,i,j,k])):
           osx=omega_yz[h,i,j,k]
        else:
           continue
        if(math.isfinite(omega_zx[h,i,j,k])):
           osy=omega_zx[h,i,j,k]
        else:
           continue
        if(math.isfinite(omega_xy[h,i,j,k])):
           osz=omega_xy[h,i,j,k]
        else:
           continue
        #we compute S in the lab frame
        fac=1/(4*m)
        Sx=fac*(Ep*osx+(py*otz-pz*oty))
        Sy=fac*(Ep*osy+(pz*otx-px*otz))
        Sz=fac*(Ep*osz+(px*oty-py*otx))
        #we compute S in the particle LRF frame
        bof=(px*Sx+py*Sy+pz*Sz)/(Ep*(m+Ep))
        Sx_lrf=Sx-bof*px
        Sy_lrf=Sy-bof*py
        Sz_lrf=Sz-bof*pz
 
        datas[index,0:4]=indata[0:4] #we copy the coordinates t,x,y,z
        datas[index,4:6]=pt,rapidity
        datas[index,6:9]=Sx_lrf,Sy_lrf,Sz_lrf

        index=index+1

    if(count_reads > nlines):
        print("Hey, something here went wrong... I counted "+str(count_lines)+" read in a block of "+hadfile+", but they should have been at most "+str(nlines))
        print("Please, check...")
        sys.exit(3)

    for a in range(index):
        outfile.write('{:7.3f}'.format(datas[a,0])+sp+'{:7.3f}'.format(datas[a,1])+sp+'{:7.3f}'.format(datas[a,2])+sp+'{:7.3f}'.format(datas[a,3])+sp+'{:7.3f}'.format(datas[a,4])+sp+'{:7.3f}'.format(datas[a,5]))
        for q in range(6,9):
            outfile.write(sp+'{:12.8e}'.format(datas[a,q]))
        outfile.write("\n")

    if(count_reads < nlines):
        break
    else:
        count_lines=count_lines+nlines
        count_reads=0
        index=0

hf.close()
outfile.close()
print("All done!")