aboutsummaryrefslogtreecommitdiff
path: root/utils/plot-fit-results
AgeCommit message (Collapse)Author
2019-07-12set ignore_index=True in pd.concat() since we don't care about the indextlatorre
2019-07-12fix a couple of bugs in plot-energytlatorre
This commit fixes two small bugs in the plotting scripts. First, after the HDF5 commit I wasn't correctly computing the particle ID string which I had been using before which was needed in order to plot things correctly. Second, I realized that the dataframe groupby function first() actually selects the first non-null column from each group! What I really wanted was the first row from each group, so all instances of .first() were updated to .nth(0). See https://stackoverflow.com/questions/20067636/pandas-dataframe-get-first-row-of-each-group.
2019-07-11switch from YAML output to HDF5 to speed things uptlatorre
2019-06-05try to import CLoader if possible since it's *much* fastertlatorre
2019-03-26update plotting scripts to handle case when there is no fittlatorre
2019-03-16add GPLv3 licensetlatorre
2019-01-15update zebra library to be able to use linkstlatorre
This commit updates the zebra library files zebra.{c,h} so that it's now possible to traverse the data structure using links! This was originally motivated by wanting to figure out which MC particles were generated from the MCGN bank (from which it's only possible to access the tracks and vertices using structural links). I've also added a new test to test-zebra which checks the consistency of all of the next/up/orig, structural, and reference links in a zebra file.
2018-11-30update plot-fit-results to handle masked arrays properlytlatorre
2018-11-25add a script to plot the fit results as a function of energytlatorre
d='n150' href='#n150'>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 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414
#!/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
from itertools import islice

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_ball(R):
    """
    Generates a random point inside a sphere of radius R.
    """
    while True:
        pos = (np.random.rand(3)*2-1)*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`. Returns the 4-vector as would
    be seen by a frame moving 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_gen_decay1():
    """
    A super simple test that if we generate a decay with E = mass, the two
    particles come out back to back.
    """
    for v1, v2 in islice(gen_decay(100,100,0.5,0.5),100):
        dir1 = v1[1:]/np.linalg.norm(v1[1:])
        dir2 = v2[1:]/np.linalg.norm(v2[1:])
        assert np.isclose(np.dot(dir1,dir2),-1.0)

def test_gen_decay2():
    """
    A super simple test that if we generate a decay with E >> mass, the two
    particles come out in the same direction.
    """
    for v1, v2 in islice(gen_decay(100,100e3,0.5,0.5),100):
        dir1 = v1[1:]/np.linalg.norm(v1[1:])
        dir2 = v2[1:]/np.linalg.norm(v2[1:])
        assert np.isclose(np.dot(dir1,dir2),1.0,atol=1e-3)

def test_lorentz_boost2():
    """
    Simple test of lorentz_boost() using a problem from
    http://electron6.phys.utk.edu/PhysicsProblems/Mechanics/8-Relativity/decay%20massless.html.
    """
    M = 140.0
    P = 2e3
    E = np.sqrt(P**2 + M**2)
    m1 = 105.0
    m2 = 0.0
    # muon in lab frame is going in +z direction which means that the direction
    # of the muon in the pion decay frame is also in the +z direction
    v = np.array([0,0,1])
    # 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)

    if not np.isclose(E1,109.375):
        print("Energy of muon in pion rest frame = %.2e but expected %.2e" % (E1,109.375))

    assert np.isclose(E1,109.375)

    E2 = np.sqrt(m2**2 + p1**2)
    v1 = [E1] + (p1*v).tolist()
    v2 = [E2] + (-p1*v).tolist()
    # pion in lab frame is going in the +z direction, therefore the lab frame
    # is going in the -z direction relative to the pion frame
    n = np.array([0,0,-1])
    beta = np.sqrt(E**2-M**2)/E

    if not np.isclose(beta,0.99756,rtol=1e-3):
        print("Speed of pion frame relative to lab frame = %.5f but expected %.5f" % (beta,0.99756))

    assert np.isclose(beta,0.99756,rtol=1e-3)

    v1_new = lorentz_boost(v1,n,beta)
    v2_new = lorentz_boost(v2,n,beta)

    if not np.isclose(v1_new[0],2003.8,atol=1e-2):
        print("Energy of muon in lab frame = %.1f but expected %.1f" % (v1_new[0],2003.8))

    assert np.isclose(v1_new[0],2003.8,atol=1e-2)

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
    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=None,
                        help="total energy of mediator")
    parser.add_argument("-T", type=float, default=0.0,
                        help="kinetic 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 is not None:
        E = args.E
    else:
        E = args.T + args.M

    if 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,E,m1,m2),args.n):
            pos = rand_ball(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)))