aboutsummaryrefslogtreecommitdiff
path: root/utils/gen-dark-matter
blob: 007dde309ab00dc49d983bede01102a300cc9946 (plain)
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
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
#!/usr/bin/env python
# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.
"""
Script to produce MCPL files for simulating self destructing dark matter. The
arguments to the script are the dark matter mediator mass and energy and the
particle IDs of the two decay products. For example, to simulate a dark matter
mediator with a mass of 100 MeV and total energy 1 GeV decaying to an electron
and a positron:

    $ ./gen-dark-matter -M 100.0 -E 1000.0 -p1 20 -p2 21 -o output

which will create the directory "output_[UUID]" and produce three files:

    1. A MCPL file output.mcpl
    2. A MCPL header file called mcpl_header.dat
    3. A SNOMAN command file which can be used to simulate the dark matter

To simulate it with SNOMAN you can run:

    $ cd output_[UUID]
    $ snoman.exe -c output.cmd
"""

from __future__ import print_function, division
import numpy as np
import string
import uuid

SNOMAN_MASS = {
    20: 0.511,
    21: 0.511,
    22: 105.658,
    23: 105.658
}

PSUP_RADIUS = 840.0 # cm

# generate a UUID to append to all the filenames so that if we run the same job
# twice we don't overwrite the first job
ID = uuid.uuid1()

SNOMAN_TEMPLATE = \
"""
$processor_list 'MCO UCL PRU PCK OUT END'

* set up output file
$output_format $full_ds
$zdab_option $zdab_max_mc
file out 1 @output
* Don't know if this next line is necessary
file MCO 1 ./MC_Atm_Nu_No_Osc_Snoman_Genie10000_X_0_rseed.dat checkpoint=10
file mco 2 @mcpl
$mc_event_rate -0.003189 $per_sec


$prune_mc               $keep
$prune_mcpm             $keep
$prune_mcvx_source      $keep
$prune_mcvx_boundary    $drop
$prune_mcvx_interaction $drop
$prune_mcvx_sink        $drop
$prune_mcvx_pre_source  $drop

$mcrun 10000

* simulate run conditions for run 10000
$mc_gen_run_cond $on

* Don't think this is necessary since it's reset in run_mc_atmospherics
$num_events @num_events

* titles files for run 10000
titles /sno/mcprod/dqxx/DQXX_0000010000.dat
titles /sno/mcprod/anxx/titles/pca/anxx_pca_0000010623_p2.dat
titles /sno/output/anxx/titles/neutrino/10000-10999/anxx_nu_0000010000_p12.dat

* set the average number of noise hits per event
* this comes from the autosno generated MC_Atm_Nu_No_Osc_Snoman_Genie10000_X_1_run_mc_atm_nu_genie.cmd file
set bank TRSP 1 word 2 to 2.051309

* From activate_atmospherics.cmd

$killvx                 7
$killvx_neutron		5
$egs4_ds		$off
$store_full_limit	10
$max_cer_ge_errors      2000

* Enable hadron propagation

titles sno_hadron_list.dat
titles chetc_sno.dat
titles flukaaf_sno.dat
$enable_hadrons         $on

* Load information for muons

titles music_sno_info.dat
titles music_double_diff_rock.dat
titles muon.dat
titles muon_param.dat
titles photo_dis.dat
$enable_music_calc      $off

@load_d2o_settings.cmd
@run_mc_atmospherics
""".strip()

MCPL_HEADER="""
*DO  MCPL 1 -i(30I 2I / 2I 17F) -n2000000
#.
#.  This bank is used to run particles through SNOMAN (with file MCO 2 ... command)
#.  This file was derived from a .dat file of Christian Nally.
#.
#.    Contact:  D. Waller (Carleton)
#.
#.       Standard Database Header
#.
19750101        0 20380517 03331900  #.  1..4   Intrinsic validity
       0        0        0           #.  5..7   Data type, Task type, Format no.
       0        0        0           #.  8..10  Creation Date, Time, Source Id.
19750101        0 20380517 03331900  #. 11..14  Effective validity
       0        0                    #. 15..16  Entry Date Time
4*0                                  #. 17..20  Spare
1000000*0                            #. 21..30  Temporary data (not in database)
#.
#.    End of Standard Database Header
#.
#.    User Data.
"""

class MyTemplate(string.Template):
    delimiter = '@'

def rand_sphere2(R):
    """
    Generates a random point inside a sphere of radius R.
    """
    while True:
        pos = np.random.rand(3)*R

        if np.linalg.norm(pos) < R:
            break

    return pos

def rand_sphere():
    """
    Generates a random point on the unit sphere.
    """
    u = np.random.rand()
    v = np.random.rand()

    phi = 2*np.pi*u
    theta = np.arccos(2*v-1)

    dir = np.empty(3,dtype=float)

    dir[0] = np.sin(theta)*np.cos(phi)
    dir[1] = np.sin(theta)*np.sin(phi)
    dir[2] = np.cos(theta)

    return dir

def lorentz_boost(x,n,beta):
    """
    Performs a Lorentz boost on the 4-vector `x` along the direction `n` at a
    velocity `beta` (in units of the speed of light).

    Formula comes from 46.1 in the PDG "Kinematics" Review. See
    http://pdg.lbl.gov/2014/reviews/rpp2014-rev-kinematics.pdf.
    """
    n = np.asarray(n,dtype=float)
    x = np.asarray(x,dtype=float)

    gamma = 1/np.sqrt(1-beta**2)

    # normalize direction
    n /= np.linalg.norm(n)

    # get magnitude of `x` parallel with the boost direction
    x_parallel = np.dot(x[1:],n)

    # compute the components of `x` perpendicular
    x_perpendicular = x[1:] - n*x_parallel

    E = x[0]
    E_new = gamma*E - gamma*beta*x_parallel
    x_new = -gamma*beta*E + gamma*x_parallel

    x_new = [E_new] + (x_new*n + x_perpendicular).tolist()

    return np.array(x_new)

def gen_decay(M, E, m1, m2):
    """
    Generator yielding a tuple (v1,v2) of 4-vectors for the daughter particles
    of a 2-body decay from a massive particle M with total energy E in the lab
    frame.

    Arguments:

            M  - mass of parent particle (MeV)
            E  - Total energy of the parent particle (MeV)
            m1 - mass of first decay product
            m2 - mass of second decay product
    """
    while True:
        # direction of 1st particle in mediator decay frame
        v = rand_sphere()
        # calculate momentum of each daughter particle
        # FIXME: need to double check my math
        p1 = np.sqrt(((M**2-m1**2-m2**2)**2-4*m1**2*m2**2)/(4*M**2))
        E1 = np.sqrt(m1**2 + p1**2)
        E2 = np.sqrt(m2**2 + p1**2)
        v1 = [E1] + (p1*v).tolist()
        v2 = [E2] + (-p1*v).tolist()
        # random direction of mediator in lab frame
        n = rand_sphere()
        beta = np.sqrt(E**2-M**2)/E
        v1_new = lorentz_boost(v1,n,beta)
        v2_new = lorentz_boost(v2,n,beta)

        yield v1_new, v2_new

def test_lorentz_boost():
    """
    Super simple test of the lorentz_boost() function to make sure that if we
    boost forward by a speed equal to the particle's original speed, it's
    4-vector looks like [mass,0,0,0].
    """
    m1 = 100.0
    beta = 0.5
    p1 = m1*beta/np.sqrt(1-beta**2)
    p1 = [np.sqrt(m1**2 + p1**2)] + [p1,0,0]

    p1_new = lorentz_boost(p1,[1,0,0],0.5)

    assert np.allclose(p1_new,[m1,0,0,0])

if __name__ == '__main__':
    import argparse
    from itertools import islice
    import sys
    import os

    parser = argparse.ArgumentParser("generate MCPL files for self destructing dark matter")
    parser.add_argument("-M", type=float, default=100.0,
                        help="mass of mediator")
    parser.add_argument("-E", type=float, default=100.0,
                        help="total energy of mediator")
    parser.add_argument("-p1", type=int, default=20,
                        help="SNOMAN particle ID for 1st decay product")
    parser.add_argument("-p2", type=int, default=21,
                        help="SNOMAN particle ID for 2nd decay product")
    parser.add_argument("-n", type=int, default=100,
                        help="number of events to generate")
    parser.add_argument("-o", "--output", type=str, required=True,
                        help="output prefix")
    args = parser.parse_args()

    if args.p1 not in SNOMAN_MASS:
        print("%i is not a valid particle ID" % args.p1,file=sys.stderr)
        sys.exit(1)

    m1 = SNOMAN_MASS[args.p1]

    if args.p2 not in SNOMAN_MASS:
        print("%i is not a valid particle ID" % args.p2,file=sys.stderr)
        sys.exit(1)

    m2 = SNOMAN_MASS[args.p2]

    if args.M < m1 + m2:
        print("mediator mass must be greater than sum of decay product masses",file=sys.stderr)
        sys.exit(1)

    if args.E < args.M:
        print("mediator energy must be greater than or equal to the mass",file=sys.stderr)
        sys.exit(1)

    if args.n < 0:
        print("number of events must be positive",file=sys.stderr)
        sys.exit(1)

    # The format for the MCPL files as read in by SNOMAN is:
    #
    #     Word    Type     Description
    #     ----    ----     --------------------------------
    #       1      I       Number of events in the list.
    #       2      I       Number of words per track. (This may vary because
    #                      polarisation is included only for Cerenkov photons).
    #      j+1     I       Number of particles in this event.     
    #      j+2     I       Particle type of this particle.
    #      j+3   F,F,F     X,Y,Z of particle.
    #      j+6     F       Time of particle.
    #      j+7     F       Energy of particle.
    #      j+8   F,F,F     U,V,W of particle.
    #      j+11  F,F,F     x,y,z components of the particle's polarisation.

    
    new_dir = "%s_%s" % (args.output,ID.hex)

    os.mkdir(new_dir)
    os.chdir(new_dir)

    mcpl_filename = args.output + ".mcpl"

    with open(mcpl_filename, "w") as f:
        f.write("%i %i\n" % (args.n, 10))

        for v1, v2 in islice(gen_decay(args.M,args.E,m1,m2),args.n):
            pos = rand_sphere2(PSUP_RADIUS)
            p1 = np.linalg.norm(v1[1:])
            p2 = np.linalg.norm(v2[1:])
            f.write("    %i %i %f %f %f %f %f %f %f %f\n" % (2,args.p1,pos[0], pos[1], pos[2],0.0, v1[0], v1[1]/p1, v1[2]/p1, v1[3]/p1))
            f.write("    %i %i %f %f %f %f %f %f %f %f\n" % (2,args.p2,pos[0], pos[1], pos[2],0.0, v2[0], v2[1]/p2, v2[2]/p2, v2[3]/p2))

    # Write out mcpl_header.dat which is needed by the cmd file
    with open("mcpl_header.dat", "w") as f:
        f.write(MCPL_HEADER)

    template = MyTemplate(SNOMAN_TEMPLATE)

    mcds_filename = args.output + ".mcds"
    cmd_filename = args.output + ".cmd"

    with open(cmd_filename, "w") as f:
        f.write(template.safe_substitute(output=mcds_filename, mcpl=mcpl_filename, num_events=str(args.n)))
ss="p">("Log file '%s' doesn't exist. Assuming job failed." % log_file) return 7 hdf5_file = "%s.hdf5" % root try: with h5py.File(hdf5_file) as f: if 'git_sha1' in f.attrs: # Job completed successfully return 8 else: log.warn("No git_sha1 attribute in HDF5 file '%s'. Assuming job failed." % hdf5_file) return 7 except IOError: log.warn("HDF5 file '%s' doesn't exist. Assuming job failed." % hdf5_file) return 7 return 7 def get_njobs(): """ Returns the total number of jobs in the job queue. """ log.debug("condor_q -json") output = check_output(["condor_q","-json"]) if not output: return 0 data = json.loads(output) return len(data) def main(conn, dqxx_dir, max_retries, max_jobs): c = conn.cursor() results = c.execute('SELECT id, filename, run, uuid, gtid, particle_id, max_time, state, nretry, submit_file, priority FROM state ORDER BY priority DESC, timestamp ASC') njobs = get_njobs() stats = {} output = check_output(["condor_q","-json","--attributes","SubmitFile,JobStatus"]) if output: data = json.loads(output) else: data = [] for id, filename, run, uuid, gtid, particle_id, max_time, state, nretry, submit_file, priority in results.fetchall(): if state not in stats: stats[state] = 1 else: stats[state] += 1 if state == 'NEW': if njobs >= max_jobs: log.verbose("Skipping job %i because there are already %i jobs in the queue" % (id,njobs)) continue log.verbose("Creating submit file for %s gtid %i particle combo %i" % (filename, gtid, particle_id)) submit_file = create_submit_file(filename, uuid, run, gtid, dir, dqxx_dir, particle_id, max_time, priority) if submit_job(submit_file): log.warn("Failed to submit job %i: %s" % (id,str(e))) c.execute("UPDATE state SET state = 'FAILED', message = ? WHERE id = ?", ("failed to submit job: %s" % str(e),id)) else: log.notice("Successfully submitted job %i for %s gtid %i particle combo %i" % (id, filename, gtid, particle_id)) njobs += 1 c.execute("UPDATE state SET state = 'RUNNING', submit_file = ?, nretry = ? WHERE id = ?", (abspath(submit_file),0,id)) elif state == 'RUNNING': # check to see if it's completed job_status = get_job_status(submit_file, data=data) if job_status in (0,1,2,4): # nothing to do! log.verbose("Still waiting for job %i to finish" % id) elif job_status == 8: # Success! log.notice("Job %i completed successfully!" % id) c.execute("UPDATE state SET state = 'SUCCESS' WHERE id = ?", (id,)) elif job_status == 5: if nretry < max_retries: log.notice("Releasing job %i" % id) if release_job(submit_file) != 0: log.warn("Failed to release job %i. Setting it to FAILED state." % id) c.execute("UPDATE state SET state = 'FAILED', message = ? WHERE id = ?", ("failed to release job",id)) else: log.verbose("Job %i has now been retried %i times" % (id,nretry+1)) c.execute("UPDATE state SET nretry = nretry + 1 WHERE id = ?", (id,)) else: log.warn("Job %i has failed %i times. Clearing it from the queue. Setting it to FAILED state." % (id,nretry)) remove_job(submit_file) c.execute("UPDATE state SET state = 'FAILED', message = ? WHERE id = ?", ("failed too many times", id)) elif job_status == 7: # Retry if nretry < max_retries if nretry < max_retries: if submit_job(submit_file): log.warn("Failed to resubmit job %i. Setting it to FAILED state." % id) c.execute("UPDATE state SET state = 'FAILED', message = ? WHERE id = ?", ("failed to resubmit job",id)) else: log.notice("Resubmitted job %i" % id) njobs += 1 c.execute("UPDATE state SET state = 'RUNNING', nretry = ? WHERE id = ?", (nretry+1,id)) else: log.warn("Job %i has failed %i times. Setting it to FAILED state." % (id,nretry)) c.execute("UPDATE state SET state = 'FAILED', message = ? WHERE id = ?", ("failed too many times", id)) else: # Don't know what to do here for Removed or Submission_err log.warn("Job %i is in the state %i. Don't know what to do." % (id, job_status)) elif state == 'RETRY': if njobs >= max_jobs: log.verbose("Skipping job %i because there are already %i jobs in the queue" % (id,njobs)) continue log.notice("Resubmitting job %i from RETRY state" % id) if submit_job(submit_file): log.warn("Failed to resubmit job %i. Setting it to FAILED state." % id) c.execute("UPDATE state SET state = 'FAILED', message = ? WHERE id = ?", ("failed to resubmit job",id)) else: log.notice("Resubmitted job %i" % id) njobs += 1 c.execute("UPDATE state SET state = 'RUNNING', message = 'resubmitted from RETRY state', nretry = ? WHERE id = ?", (nretry+1,id)) elif state in ('SUCCESS','FAILED'): # Nothing to do here pass else: log.warn("Job %i is in the unknown state '%s'." % (id,state)) conn.commit() log.notice("Stats on jobs in the database:") for state, value in stats.iteritems(): log.notice(" %s: %i" % (state,value)) def array_to_particle_combo(combo): particle_combo = 0 for i, id in enumerate(combo[::-1]): particle_combo += id*100**i return particle_combo if __name__ == '__main__': import argparse from subprocess import check_call import os import tempfile import h5py from itertools import combinations_with_replacement import sqlite3 import traceback from sddm.plot_energy import prompt_event, gtid_sort, unwrap_50_mhz_clock import pandas as pd from sddm.dc import DC_MUON, DC_JUNK, DC_CRATE_ISOTROPY, DC_QVNHIT, DC_NECK, DC_FLASHER, DC_ESUM, DC_OWL, DC_OWL_TRIGGER, DC_FTS, DC_ITC, DC_BREAKDOWN from sddm import read_hdf from sddm.plot_energy import get_events parser = argparse.ArgumentParser("submit grid jobs", formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("filenames", nargs='*', help="input files") parser.add_argument("--min-nhit", type=int, help="minimum nhit to fit an event", default=100) parser.add_argument("--max-particles", type=int, help="maximum number of particles to fit for", default=2) parser.add_argument("--skip-second-event", action='store_true', help="only fit the first event after a MAST bank", default=False) parser.add_argument("--max-time", type=float, default=3600*12, help="maximum time for fit") parser.add_argument("--db", type=str, help="database file", default=None) parser.add_argument('--loglevel', help="logging level (debug, verbose, notice, warning)", default='notice') parser.add_argument('--logfile', default=None, help="filename for log file") parser.add_argument('--max-retries', type=int, default=2, help="maximum number of times to try and resubmit a grid job") parser.add_argument('--auto', action='store_true', default=False, help="automatically loop over database entries and submit grid jobs") parser.add_argument('--max-jobs', type=int, default=100, help="maximum number of jobs in the grid queue at any time") parser.add_argument('-r','--reprocess', action='store_true', default=False, help="force reprocessing of runs which are already in the database") parser.add_argument("--data", action='store_true', default=False, help="zdab is not MC data") parser.add_argument("-p", "--priority", type=int, default=1, help="job priority") args = parser.parse_args() log.set_verbosity(args.loglevel) if args.logfile: log.set_logfile(args.logfile) if args.db is None: home = os.path.expanduser("~") args.db = join(home,'state.db') conn = sqlite3.connect(args.db) c = conn.cursor() # Create the database table if it doesn't exist c.execute("CREATE TABLE IF NOT EXISTS state (" "id INTEGER PRIMARY KEY, " "filename TEXT NOT NULL, " "run INTEGER NOT NULL, " "timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, " "uuid TEXT NOT NULL, " "gtid INTEGER NOT NULL, " "particle_id INTEGER NOT NULL, " "state TEXT NOT NULL, " "nretry INTEGER," "submit_file TEXT," "max_time REAL NOT NULL," "message TEXT," "priority INTEGER DEFAULT 1)" ) conn.commit() results = c.execute('SELECT DISTINCT filename FROM state') unique_filenames = [row[0] for row in results.fetchall()] if 'SDDM_DATA' not in os.environ: log.warn("Please set the SDDM_DATA environment variable to point to the fitter source code location", file=sys.stderr) sys.exit(1) dir = os.environ['SDDM_DATA'] if 'DQXX_DIR' not in os.environ: log.warn("Please set the DQXX_DIR environment variable to point to the directory with the DQXX files", file=sys.stderr) sys.exit(1) dqxx_dir = os.environ['DQXX_DIR'] # get absolute paths since we are going to create a new directory dir = abspath(dir) dqxx_dir = abspath(dqxx_dir) zdab_cat = which("zdab-cat") if zdab_cat is None: log.warn("Couldn't find zdab-cat in path!",file=sys.stderr) sys.exit(1) if args.auto: try: main(conn,dqxx_dir,args.max_retries,args.max_jobs) conn.commit() conn.close() except Exception as e: log.warn(traceback.format_exc()) sys.exit(1) sys.exit(0) # generate a UUID to append to all the filenames so that if we run the same job # twice we don't overwrite the first job ID = uuid.uuid1() for filename in args.filenames: log.notice("Analyzing file %s" % filename) filename = abspath(filename) if os.path.getsize(filename)/1e6 > 500: log.warn("Skipping %s because the file size is %i MB!" % (filename, os.path.getsize(filename)/1e6)) continue with open(os.devnull, 'w') as f: # Create a temporary file to store the events. Docs recommended # this method instead of mkstemp(), but I think they're the same. output = tempfile.NamedTemporaryFile(suffix='.hdf5',delete=False) output.close() if args.skip_second_event: check_call([zdab_cat,"--skip-second-event",filename,"-o",output.name],stderr=f) else: check_call([zdab_cat,filename,"-o",output.name],stderr=f) ev, fits = get_events([output.name], merge_fits=False) if len(ev) == 0: continue if len(ev) and not args.reprocess and filename in unique_filenames: head, tail = split(filename) log.notice("Skipping %s because it's already in the database" % tail) continue head, tail = split(filename) if 'muon' in tail: ev['muon'] = True nevents = 0 njobs = 0 for index, event in ev.iterrows(): if event['nhit_cal'] >= args.min_nhit: if event['muon'] and event['stopping_muon']: pass else: if event['prompt']: if event['flasher'] or event['muon'] or event['neck']: # Only want to submit 10% of prompt flasher, muon, # and neck events if event['gtid'] % 10 != 0: continue if not event['prompt'] and event['instrumental']: # Only submit followers if they have no data cleaning cuts continue nevents += 1 for i in range(1,args.max_particles+1): for particle_combo in map(array_to_particle_combo,combinations_with_replacement([20,22],i)): c.execute("INSERT INTO state (" "filename , " "run , " "uuid , " "gtid , " "particle_id , " "max_time , " "state , " "nretry , " "priority ) " "VALUES (?, ?, ?, ?, ?, ?, 'NEW', NULL, ?)", (filename,int(event['run']),ID.hex,int(event['gtid']),particle_combo,args.max_time,args.priority)) njobs += 1 conn.commit() log.notice("submitted %i jobs for %i events" % (njobs, nevents)) # Delete temporary HDF5 file os.unlink(output.name) conn.close()