DR_functions.py 10.7 KB
Newer Older
Olga Bessidskaia Bylund's avatar
Olga Bessidskaia Bylund committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
import shutil
import os
import re
import glob

#General functions

def make_support_file(infile,outfile,test_string):
    with open(outfile,"w") as myfile:
        with open(infile,'r') as f:
            for line in f:
                if test_string in line:
                    myfile.write(line)


def return_out_index(process): #which index does the final state particle have
        p_in,p_out = process.split(" > ")
        p_in = p_in.split(" ")
        p_out = p_out.split(" ")
        particles = p_in+p_out

        windex = -1
        bindex = -1
        for i,particle in enumerate(particles):
            if (particle in (["w-","w+"])) and i>=len(p_in): #is it a FS particle?
                windex = i+1
            elif (particle in (["b","b~"])) and i>=len(p_in):
                bindex = i+1


        return bindex, windex

def find_matrix_files(P_folders):
    mfiles = []
    for mydir in P_folders:
        m = glob.glob(mydir+"/matrix_*.f")
        for i in range(len(m)):
            mfiles.append(m[i])
    return mfiles

# Finds the amplitudes that overlap with tt(+X)
def find_W_prepare_DRXhack(mfile,bindex,windex,redefine_twidth,which_DR):
    W_vector=[]
    updated_vector=[]
    the_W=""
    with open(mfile,"r") as f, open("mytmp.txt","w") as tmpf:
        to_replace={}
        mylines = f.readlines()
        for i,line in enumerate(mylines):
            if "1,"+str(bindex)+"))" in line.split("(")[-1] and "XXXXX" not in mylines[i]: # is b quark being rewritten? Beware if >10
                bindex=-1
            elif "1,"+str(windex)+"))" in line.split("(")[-1]  and "XXXXX" not in mylines[i]: # is w boson being rewritten? Beware if >10
                windex=-1

            if W_vector!=[]:
                for j in range(len(W_vector)):
                    if W_vector[j] in line and "AMP(" in line: #save AMP
                        amplitude = line.split(",")[-1]
                        if which_DR==1:
                            tmpf.write("      "+amplitude[:-2]+"=(0d0,0d0)\n") #set to zero (DR)
                        elif which_DR==2:
                            tmpf.write(amplitude[:-2]+"\n") #save amplitude
                    elif W_vector[j] in line and "AMP" not in line:
                        if W_vector[j]+")" in line: # W is redefined, stop keeping track of it
                            updated_vector.remove(W_vector[j])
                        else: # a new W is defined, containing a flagged W. keep track of it
                            the_W = line.split(",")[-2]+","+line.split(",")[-1][0:-2]
                            #print(the_W)
                            updated_vector.append(the_W)
                W_vector=updated_vector


            if "W(1,"+str(bindex)+")" in mylines[i] and "W(1,"+str(windex)+")" in mylines[i] and "MDL_WT" in mylines[i]:
                the_W = line.split(",")[-2]+","+line.split(",")[-1][0:-2]
                #print("Adding new W: ", the_W)
                W_vector.append(the_W)
                updated_vector=W_vector

                if redefine_twidth:
                    edited_line = line.replace("MDL_WT","MDL_WT_nonzero")
                    to_replace[i]=edited_line
    return to_replace
#-------------------------------------

def do_DR1_hacks(mfile,tmpfile):
    shutil.copyfile(mfile, mfile+"~" )
    with open(mfile,"w") as destination, open(mfile+"~") as source, open(tmpfile,"r") as mytmp:
        print "Performing MadSpin hack for ", mfile
        for line in source:
            if "JAMP(1)=" in line:
                print("performing DR1 for file ", mfile)
                destination.write("\nC     DR hack\n")
                for hackline in mytmp:
                    destination.write(hackline)
                    print(hackline)
                destination.write("C     End DR hack. \n\n")
            destination.write( line)

#-------------------------------
#DR2 helper functions


def do_driver_hacks(driver_file): #takes care of negative weights in code for driver.f for madspin
        shutil.copyfile(driver_file, driver_file+"~" )

        test_string1 = "weight.gt."
        repl_string1 = "abs(weight).gt."

        test_string2 = "maxweight=weight"
        repl_string2 = "maxweight=abs(weight)"

        with open(driver_file,"w") as destination, open(driver_file+"~","r") as source:
                mylines = source.readlines()
                for i,line in enumerate(mylines):
                        if test_string1 in line:
                                line = line.replace(test_string1,repl_string1)
                        if test_string2 in line:
                                line = line.replace(test_string2,repl_string2)
                        destination.write(line)

        os.remove(driver_file+"~")

def do_fks_hacks(pdir): #takes care of negative weights for event generation code
        shutil.copyfile(pdir+"fks_singular.f", pdir+"fks_singular.f~" )

        test_string = "wgt.lt.0.d0"
        end_string = "return"

        with open(pdir+"fks_singular.f","w") as destination, open(pdir+"fks_singular.f~") as source:
                mylines = source.readlines()
                dont_comment = True
                for i,line in enumerate(mylines):
                        if dont_comment:
                                if test_string in line:
                                        dont_comment=False
                                        destination.write("C"+line)
                                else:
                                        destination.write(line)

                        else:
                                if end_string in line:
                                        dont_comment=True
                                        destination.write(line)
                                else:
                                        destination.write("C"+line)

        os.remove(pdir+"fks_singular.f~")

def do_coupl_hacks(pdir): #define the width of the top quark for the event generation code
        shutil.copyfile(pdir+"coupl.inc", pdir+"coupl.inc~" )

        test_string1 = "MDL_WT"
        test_string2 = "WIDTHS"


        with open(pdir+"../Cards/param_card.dat","r") as pcard: ###XXX in Cards
                the_lines = pcard.readlines()
                for line in the_lines:
                        if "DECAY" in line and " 6 " in line:
                                relevant_line = line

        top_width =  re.compile("(\d.\d+e\+\d\d)").search(relevant_line).group(1)

        with open(pdir+"coupl.inc","w") as destination, open(pdir+"coupl.inc~") as source:
                mylines = source.readlines()
                for i,line in enumerate(mylines):
                        destination.write(line)
                        if test_string1 in line and test_string2 in line:
                                destination.write("\n      DOUBLE PRECISION MDL_WT_nonzero\n")
                                destination.write("      PARAMETER (MDL_WT_nonzero="+top_width+")\n")

        os.remove(pdir+"coupl.inc~")
#-------------------------
# DR2 specific

def find_jamp(mfile,helperfile,test_string,end_string):

        with open(mfile,"r") as source:
                mylines = source.readlines()
                dont_store = True
                my_jamp = ""
                for i,line in enumerate(mylines):
                        if dont_store:
                                if test_string in line:
                                        dont_store=False
                                        my_jamp += line

                        else:
                                if end_string in line:
                                        dont_store=True
                                else:
                                        my_jamp += line

        with open(helperfile,"w") as destination:
                destination.write(my_jamp)
                print "my_jamp = ", my_jamp


def do_DR2_hack(mfile, tmpfile, suffix,to_replace,should_replace): #to_replace and should_replace refer to whether to put top propagator on shell - do so for madgraph code but not for madspin
    shutil.copyfile(mfile, mfile+"~" )
    with open(mfile,"w") as destination, open(mfile+"~") as source, open(tmpfile,"r") as mytmp, open(mfile+".jamp1","r") as saved_jamp1, open(mfile+".jamp2","r") as saved_jamp2:
        tmp_content = mytmp.readlines()
        for i,line in enumerate(source):
            if i in to_replace.keys() and should_replace==1:
                destination.write(to_replace[i])
            else:
                destination.write(line)

            if "JAMP(NCOLOR)" in line and "COMPLEX*16" in line:
                destination.write("      REAL*8 MATRIX"+suffix+"_res\n")
                destination.write("      COMPLEX*16 JAMP_res(NCOLOR)\n")
            elif "JAMP(1)=" in line:
                print("performing DR2 for file ", mfile)
                jamplines = saved_jamp1.readlines()
                my_jamp1 = ""
                for j,jline in enumerate(jamplines):
                        my_jamp1 +=  jline
                jamp_res1 = "JAMP_res(1)="+write_jamp_res(my_jamp1,tmp_content)
                #jamp_prefix = line.split("=")[1].split("*(")
            elif "JAMP(2)=" in line:
                jamplines = saved_jamp2.readlines()
                my_jamp2 = ""
                for j,jline in enumerate(jamplines):
                        my_jamp2 +=  jline

                jamp_res2 = "JAMP_res(2)="+write_jamp_res(my_jamp2,tmp_content)

            elif "DENOM(I)" in line:
                line=next(source)
                destination.write(line)
                destination.write("\n      "+jamp_res1)
                destination.write("      "+jamp_res2+"\n")
                hardcode = "      MATRIX"+suffix+"_res = 0.D0\n      DO I = 1, NCOLOR\n        ZTEMP = (0.D0,0.D0)\n        DO J = 1, NCOLOR\n          ZTEMP = ZTEMP + CF(J,I)*JAMP_res(J)\n        ENDDO\n        MATRIX"+suffix+"_res = MATRIX"+suffix+"_res+ZTEMP*DCONJG(JAMP_res(I))/DENOM(I)\n      ENDDO\n      MATRIX"+suffix+" = MATRIX"+suffix+" - MATRIX"+suffix+"_res\n"
                destination.write(hardcode)

    os.remove(mfile+"~")

def write_jamp_res(my_jamp_line,tmp_c):
    line = my_jamp_line
    my_jamp_line = my_jamp_line.replace("+",",+")
    my_jamp_line = my_jamp_line.replace("-",",-")
    #print "original jamp: ", my_jamp_line
    split_line = my_jamp_line.split(",")[1:]
    jamp_resX = ""
    for amp in tmp_c:
        for term in split_line:
            if amp[:-1] in term:
                jamp_resX = jamp_resX+term

    jamp_prefix = line.split("=")[1].split("*(")
    if len(jamp_prefix)>1:
        jamp_resX = jamp_prefix[0]+"*("+jamp_resX
        if jamp_resX[-2:]!="\n":
            if "))" not in jamp_resX[-4:]:
                jamp_resX=jamp_resX+")\n"
            else:
                jamp_resX=jamp_resX+"\n"

    else:
        if "\n" not in jamp_resX:
            jamp_resX=jamp_resX+"\n"

    return jamp_resX