#!/usr/bin/env python

"""
This script solves the problem posed in Computational Physics Lab 4, Part 2:
Calculating the Muon Lifetime.

In order to run this script, you need the utility modules v3.py, rk.py,
and quasi_random.so. You can get these by following the instructions in
the lab description.
"""

__author__="Igor Volobouev (i.volobouev@ttu.edu)"
__version__="0.1"
__date__ ="Feb 22 2008"

import math
from random import random
import sobol
from rk import *
from v3 import *
from pylab import *

def phase_space_decay(parent, m1, m2, rnd0, rnd1):
    """
    This function decays the parent particle into two daughter particles
    with masses m1 and m2. The random numbers rnd0 and rnd1 determine
    the direction of the first particle. The second particle will have
    the opposite direction in the parent's rest frame. The magnitude
    of the particle momentum is determined from energy-momentum
    conservation. The function returns the tuple of 4-momenta of decay
    products. The daughter particle with mass m1 is the first in the tuple.
    """
    # Check that the mass arguments make sense. The decay is possible
    # only if the mass of the parent is at least as large as the
    # combined mass of the decay products.
    assert m1 >= 0.0 and m2 >= 0.0
    assert parent.m >= m1 + m2
    
    # Calculate the boost which will take us from the rest frame
    # of the parent into the lab rest frame
    boost = Boost(parent).inverse()

    if parent.m == m1 + m2:
        # Special case: there is not enough energy for decay
        # products to fly apart
        p1 = P4(v3.V3(0, 0, 0), m1)
        p2 = P4(v3.V3(0, 0, 0), m2)
    else:
        # Calculate the direction of one of the particles
        cosTheta = rnd0*2.0 - 1.0
        sinTheta = math.sqrt(1.0 - cosTheta*cosTheta)
        phi = rnd1*2.0*math.pi
        # Calculate the magnitude of the momentum of the decay products
        # in the parent's rest frame
        magnitude = kinematic_lambda(parent.m*parent.m, m1*m1,
                                     m2*m2)/2.0/parent.m
        # Construct the 3-momentum vector of the first particle
        momentum = magnitude*v3.V3(sinTheta*math.cos(phi),
                                   sinTheta*math.sin(phi),
                                   cosTheta)
        # Construct the 4-momentum vector of the first particle
        p1 = P4(momentum, m1)
        # Construct the 4-momentum vector of the second particle
        p2 = P4(-momentum, m2)

    # Return the decay products boosted into the lab frame
    return boost(p1), boost(p2)


class MuonDecayME(object):
    """
    This callable returns the matrix element squared for the muon decay
    """
    def __init__(self, GFermi):
        self.GFermi = GFermi
    def __call__(self, p0, p1, p2, p3):
        return self.GFermi*self.GFermi*64.0*p0.dot(p2)*p3.dot(p1)


class ThreeBodyDecayWidth(object):
    """
    This class calculates the decay width of a parent particle into
    three daughter particles. All calculations are performed in the class
    constructor, and various class functions are used to access the
    results. The constructor arguments are as follows:
    
    matrel     -- The callable for the matrix element squared of the process
    
    mMother    -- Mass of the decaying particle
    
    m1, m2, m3 -- Masses of the daughter particles. Note that the matrix
                  element callable will be invoked using the 4-momenta
                  ordered as follows: p0, p1, p2, p3. p0 corresponds to
                  the decaying particle, p1 corresponds to the daughter
                  with mass m1, etc.
    
    precision  -- The relative precision with which the decay width will
                  be calculated. This is the ratio of the estimated standard
                  deviation of the decay width to the decay width. Note
                  that this precision estimate is probabilistic, not absolute.
                  We will satisfy this precision about 68% of the time.
                  The precision estimate works for Monte Carlo calculation
                  using pseudo-random points. It does not work for
                  quasi-random points.
                  
    maxcycles  -- The calculation stops unconditionally after using this
                  number of random points. Specify this number as 0 to set
                  the limit to the length of the list of random points.
                  
    randomSet  -- A list of 5d pseudo (or quasi) random points which will
                  be used in the calculation. Should have at least "maxcycles"
                  elements (if "maxcycles" is not 0).

    checkFreq  -- How often we will check that the precision target is
                  satisfied (makes sense for pseudo-random points only).
    """
    def __init__(self, matrel, mMother, m1, m2, m3,
                 precision, maxcycles, randomSet, checkFreq):
        # Make sure the arguments make sense
        assert mMother >= 0.0
        assert m1 >= 0.0
        assert m2 >= 0.0
        assert m3 >= 0.0
        assert mMother > m1 + m2 + m3
        assert maxcycles >= 0

        if maxcycles == 0:
            maxcycles = len(randomSet)
        if checkFreq <= 0:
            checkFreq = 1

        # Remember various quantities we will need in the calculation
        self.matrel = matrel
        self.mother = P4(V3(0, 0, 0), mMother)
        self.m1 = m1
        self.m2 = m2
        self.m3 = m3
        
        # Minimum value of compound mass squared
        self.minmsq = (m2 + m3)*(m2 + m3)
        # Maximum value of compound mass squared
        self.maxmsq = (mMother - m1)*(mMother - m1)
        # Constant factor for the integral
        self.factor = (self.maxmsq - self.minmsq)/256.0/pow(math.pi*mMother, 3)
        
        # Accumulate the statistics while we are going along:
        # the number of points, the sum of the integrand values,
        # the sum of the squares of the integrand values. These
        # are needed to calculate the mean and the standard deviation
        # of the integrand. Note that the method of calculating
        # the standard deviation as a difference between the average
        # of squares and the average squared exhibits strong subtractive
        # cancellation. We use it here only because the requested
        # precision of the integral (1%) does not present a big
        # problem.
        self.count = 0
        self.sum = 0.0
        self.sumsq = 0.0
        
        # Loop over the random points
        while self.count < maxcycles:
            w = self._weight(randomSet[self.count])
            self.sum += w
            self.sumsq += w*w
            self.count += 1
            if self.count % checkFreq == 0:
                # Have we reached the requested precision?
                # If yes then finish the calculation
                r = self.uncertainty()/self.width()
                if r < precision:
                    break

    def uncertainty(self):
        """
        Estimate the integral uncertainty using the number of points,
        the sum of integrand values, and the sum of squares of the
        integrand values.
        """
        ave = self.sum/self.count;
        var = self.sumsq/self.count - ave*ave;
        if var > 0.0:
            # Return the uncertainty of the average which equals to
            # the standard deviation of the sample divided by sqrt(N)
            return math.sqrt(var/self.count);
        else:
            return 0.0
        
    def width(self):
        """
        Returns the current estimate of the integral
        """
        return self.sum/self.count;
    
    def _weight(self, randomPoint):
        """
        Returns the integrand value which corresponds to the given
        5d random point
        """
        # Generate the compound m squared
        mcsq = self.minmsq + randomPoint[0]*(self.maxmsq - self.minmsq);
        mc = math.sqrt(mcsq)
        
        # Generate the decay of the mother into the first
        # particle and the compound
        p1, pc = phase_space_decay(self.mother, self.m1, mc,
                                   randomPoint[1], randomPoint[2])
        
        # Generate the decay of the compound
        p2, p3 = phase_space_decay(pc, self.m2, self.m3,
                                   randomPoint[3], randomPoint[4])

        # Calculate the matrix element
        me = self.matrel(self.mother, p1, p2, p3);

        # Return the integrand value
        return self.factor*me*\
            kinematic_lambda(mcsq, self.m2*self.m2, self.m3*self.m3)/mcsq*\
            kinematic_lambda(self.mother.m*self.mother.m,self.m1*self.m1,mcsq)


def muon_lifetime(relativeEps, maxN, randomSeq, checkFreq, isPseudo=0):
    """
    This function calculates the muon lifetime using the particle masses
    and the physical constants provided in the lab description
    """
    GFermi = 1.16637e-5
    mmu = 105.66e-3
    me = 0.511e-3
    hbar = 6.582119e-25
    decayME = MuonDecayME(GFermi)
    width = ThreeBodyDecayWidth(decayME, mmu, me, 0.0, 0.0,
                                relativeEps, maxN, randomSeq, checkFreq)
    w = width.width()
    tau = hbar/w

    # Print the results to the terminal
    print "Number of integrand evaluations is", width.count
    if isPseudo:
        precision = width.uncertainty()/w
        print "Muon lifetime is", tau, "+-", precision*tau, "s"
    else:
        print "Muon lifetime is", tau, "s"
    return tau


# Construct the random point sequence using the pseudo-random number
# generator which comes with the Python distribution
n_random = 65536
random_seq = []
for i in xrange(n_random):
    random_seq.append((random(), random(), random(), random(), random()))

# Calculate the muon lifetime with 1% precision using pseudo-random numbers
print "Calculation using pseudo-random sequence:"
tau_pseudo_mc = muon_lifetime(0.01, 0, random_seq, 100, 1)

# Construct the point sequence using the Sobol quasi-random sequence
sobol.init(5)
sobol_seq = []
for i in xrange(n_random):
    sobol_seq.append(sobol.next())

# Calculate the muon lifetime using the quasi-random sequence
# for different number of points in the form 2**k
klist = range(6, 17)
num_points = []
tau_quasi_mc = []
print "\nCalculation using quasi-random sequence:"
for k in klist:
    n = 2**k
    num_points.append(n)
    tau = muon_lifetime(0.0, n, sobol_seq, 1000000)
    tau_quasi_mc.append(tau)

# Plot the result as a function of k
figure()
plot(klist, tau_quasi_mc)
title('Quasi-MC Results')
xlabel('Logarithm (base 2) of the Number of Points')
ylabel('Muon Lifetime (s)')

# Assume that the result calculated last (highest k) is exact.
# Calculate the relative error w.r.t. this number.
errors = [abs(t - tau)/tau for t in tau_quasi_mc[:-1]]

# Plot the error w.r.t the number of points. When you see the plot,
# you will see that the integration precision for quasi-MC is close
# to O(1/N) instead of O(1/sqrt(N)) typical for pseudo MC.
figure()
loglog(num_points[:-1], errors)
title('Quasi-MC Convergence')
xlabel('Number of Points')
ylabel('Relative Error')

show()