DecayAnalysis

DecayAnalysis is a library implementing flexible, modularized, and optimized analysis of particle decays, integrating arbitrary high-precision floating-point arithmetic, uncertainty analysis, convenience plotting, and more.

Example Code (Updated as of 01/04/26)

Run in the following order.
Setup: First prepare the particles by specifying their names and masses.

''' Simulating tritium beta-decay'''
# Prepare particles:
tritium = ['T', MASS_TRITIUM]
helium = ['He3', MASS_HELIUM]
electron = ['e-', MASS_ELECTRON]
neutrino = ['ve', MASS_ENEUTRINO]

The name (first argument) given to each particle is important if utilizing prebuilt matrix squared element functions, which rely on knowing the physical type of particle each symbol refers to. In such cases, one should use the "conventional" name, which is easily accessible from the enumerated Particle class in particle_dict.py.

''' Class Particle in particle_dict.py'''
tritium = [Particle.TRITIUM.value, MASS_TRITIUM]
helium = [Particle.HELIUM3.value, MASS_HELIUM]
electron = [Particle.ELECTON.value, MASS_ELECTRON]
neutrino = [Particle.ELECTON_NEUTRINO.value, MASS_ENEUTRINO]

The mass (second argument) provided to each particle can be a float, or a function of the signature:

# Return array of size (n_events,). Each element in the return array must be of type Decimal (from library decimal).
def mass_fn(min_mass, max_mass, n_events) -> numpy.ndarray: 

For instance, the neutrino particle may be specified as follows:

def neutrino_mass_fn(min_mass:float, max_mass:float, n_events:int):
    masses = np.random.choice([Decimal(str(m1)), Decimal(str(m2)), Decimal(str(m3))], size=n_events, p=[p1, p2, p3])
    return masses

neutrino = ['ve', neutrino_mass_fn]

Next, specify uncertainty and smearing functions characterizing the decay. Uncertainty functions represent the experimental uncertainty (in std) characterizing the observed kinematics in a given decay. Smearing functions use the corresponding uncertainty function to simulate experiments by smearing the truth values of the simulation accordingly.

Numerous prebuilt functions are available in uncertainty.py. These functions can either be applied to a single particle's kinematic (E, px, py, pz), or it can be applied generally to a particle's momentum (p, replacing all three of px, py, pz).

''' Class Uncertainty in uncertainty.py '''
# Uncertainty functions
numerical_uncertainty = Uncertainty.numerical_decimal_uncertainty
p_numerical_uncertainty = Uncertainty.p_numerical_decimal_uncertainty

resolution_fn1 = Uncertainty.absolute_resolution_uncertainty
resolution_fn2 = Uncertainty.relative_resolution_uncertainty
p_resolution_fn = Uncertainty.p_absolute_resolution_uncertainty

# Smearing functions
gauss_smear = Uncertainty.gaussiam_smear
p_gauss_smear = Uncertainty.gaussian_angular_smear

One's own resolution function must have the following signature (see examples in uncertainty.py). Note with_prec represents whether to calculate via float64 (False) or arbitrarily-high precision mpmath.mpf (True).

# For a single kinematic, i.e. a particle's energy. Return array must be numpy.array([left_err, right_err]).
def resolution_fn(val:float|mpmath.mpf, with_prec:bool, **kwargs) -> numpy.ndarray:  

# For momentum. Return array must be a 3x2 numpy array, where each row corresponds to a different momentum component (px,py,pz)
#   and the columns are of the form [left_err, right_err].
def p_resolution_fn(three_momentum:list[float|mpmath.mpf], with_prec:bool, **kwargs) -> numpy.ndarray: 
'''If with_prec=True, *_err should be type mp.mpf (from library mpmath). If with_prec=False, *_err should be type float'''

One's own smearing function must have the following signature (see examples in uncertainty.py):

# For a single kinematic.
'''
--> weights is a numpy array of floats of shape (num_decays,), with each value chosen from a normal distribution of mean 0, std 1. This is the base randomness that may be used while smearing.
--> If with_prec=False, vals is a numpy array of shape (num_decays,), representing that kinematic across all propagated decays. If with_prec=True, it is a mpmath.matrix with num_decays rows.
--> The function returns a numpy array of shape (num_decays, 1) if with_prec=False, and a mpmath.matrix of num_decays rows if with_prec=True.
'''
def smear_fn(vals:numpy.ndarray|mpmath.matrix, res_fn:Callable, weights:numpy.ndarray, with_prec:bool, **kwargs) -> numpy.ndarray | mpmath.matrix:

# For momentum.
'''
--> weights is a list of length 3, where each item is of the form of weights for a single kinematic (representing px, py, pz, respectively)
--> If with_prec=False, three_momenta is a numpy array of shape (num_decays, 3), representing the three-momenta across all propagated decays. If with_prec=Ture, it is a mpmath.matrix with num_decays rows, 3 columns.
--> The function returns a numpy array of shape (num_decays, 3) if with_prec=False, and a mpmath.matrix of num_decays rows, 3 columns if with_prec=True.
'''
def p_smear_fn(three_momenta:numpy.ndarray|mpmath.matrix, res_fn:Callable, weights:list[numpy.ndarray], with_prec:bool, **kwargs) --> numpy.ndarray | mpmath.matrix:

Following setting or utilizing prebuilt resolution/smearing functions, one may set the matrix squared function characterizing the decay for simulation of physical decay rates.

''' Class MatrixElements in matrix_elements.py '''
# Construct/utilize matrix squared function for weighting decays.
tritium_matrix2_fn = MatrixElements.tritium_simple

One's own matrix squared function must have the following signature (see examples in matrix_elements.py):

'''
-> Receives a dictionary, where keys are particle names, and values are numpy arrays of shape (num_decays,4), with each element being of type Decimal (from library decimal). The columns represent px, py, pz, E.
-> Returns a numpy array of shape (num_decays,), with each element being of type decimal. This represents the matrix^2 element for each decay.
'''
def matrix2_fn(four_momenta_dict:dict) -> numpy.ndarray:

Configure decay:

# Configure decay and set arbitrary decimal precision.
num_decays = 10000
dp = 60

''' Class Decay in core.py '''
tritium_decay = core.Decay([tritium, helium, electron, neutrino], decimal_precision=dp, num_decays=num_decays)
# Set appropriate parent particle(s)
tritium_decay.set_parent('T', ['He3', 'e-', 've'])

''' Configure resolution/smearing. Resolution must be set for every particle/four-momentum component 
    if error propagation functions are run. '''
tritium_decay.set_resolution(['He3', 'e-', 've'], E=numerical_uncertainty, px=numerical_uncertainty, py=numerical_uncertainty, pz=numerical_uncertainty)
''' Or, more concisely, using the p_uncertainty functions and 'p' key (applies to all momentum coordinates):'''
tritium_decay.set_resolution(['He3', 'e-', 've'], E=numerical_uncertainty, p=p_numerical_uncertainty)

'''The kinematic of a particle may specify multiple resolution functions as a list, with length equal to the number of smearing functions applied. 
    Each resolution function is then applied to the corresponding smearing function.
    If there are multiple smearing functions, exactly one resolution function may still be set -- in this case, it will apply to all smearing functions.'''
# If the electron's energy passed through 2 smearing functions, one may specify 2 resolution functions for each.
tritium_decay.set_resolution(['e-'], E=[resolution_fn1, resolution_fn2], p=p_resolution_fn)


''' Smearing function does not necessarily need to be set for every particle/four-momentum component. 
    Each component may again receive a single function, or a list of multiple.''' 
# The electron's energy will be gaussian smeared twice, based on the separate resolution functions.
tritium_decay.set_smearing(['e-'], E=[gauss_smear, gauss_smear], p=p_gauss_smear)

''' One can optionally correlate the smearing functions. This smears the component listed in primary_detector, while
    reconstructing the four-momenta for the remaining coordinates to maintain sensible four-momenta. Note, the
    following overwrites our previously set smearing function for 'p', as now it is reconstructed. Unlike the above,
    this function currently works only for primary_detector='E' or 'p', not 'px', 'py', 'pz'.'''
tritium_decay.set_smearing_correlation('e-', primary_detector='E')

# Configure matrix squared function
tritium_decay.set_matrix2_fn(tritium_matrix2_fn)

''' Configure necessary kwargs for resolution/smearing functions. The values can be single values or lists of values. 
    If singular, the value will be applied to all resolution/smearing functions. Otherwise, if a list, the ith value is applied to the ith resolution/smearing function.'''
# These apply only to individual particles. In this case, resolution_fn1 gets res=1e-3, while resolution_fn2 gets res=1e-4.
tritium_decay.set_particle_kwargs('e-', res=[1e-3, 1e-4])
# These apply to all particles.
tritium_decay.set_general_kwargs(decimal_prec=dp)

# Can specify boost for parent particle in eV, natural units (c=1, as throughout). Will be utilized in next step.
zboost = 100
# Array in format [px, py, pz, E]. Should be consistent with conservation of energy.
parent_momentum = [0,0,zboost,np.sqrt(zboost**2 + MASS_TRITIUM**2)]

Finally, you may specify a pruning function, which will reject decays that do not satisfy it. Specifically, the pruning function must have the following signature:

'''
-> First argument is a dictionary, where keys are particle names, and values are numpy arrays of shape (x,4) for arbitrary x <= num_decays, with each element being of type Decimal (from library decimal). The columns represent px, py, pz, E.
-> Second argumenmt is a numpy array of shape (y,) with y <= x. The array holds the indices of the x decays which were accepted via rejection sampling by the simulation.
-> Returns a numpy array of shape (z,) with z <= y. This array prunes the accept indices of pre_accepted_indices based on the pruning criteria, returning a new accept_indices array holding the indices of the x decays which were also accepted after pruning.
'''
def prune_fn(four_momenta_dict:dict, pre_accepted_indices:numpy.ndarray) -> numpy.ndarray

The decay is now fully configured. Now time to simulate and analyze it!

Propagate decay:
Extract the four-momenta of each particle following the specified decays.

'''
Argument Specifications:
    - parent_name (string, required): name identifier of original parent particle.
    - parent_four_momentum (float array, optional): four-momentum of parent particle.
    - num_batches (integer, optional): number of batches during generation for potential speedup.
    - smearing (boolean, optional): triggers phasespace generations to be smeared. Expects smear and resolution functions to be set.
    - matrix2_weighting (boolean, optional): weights four-momenta generations considering matrix element^2. Expects matrix2 function to be set.
    - matrix2_upper_estimate_samples (integer, optional): estimates maximum matrix element^2 weight over requested domain via that number of decays, for more efficient rejection sampling.
    - multiprocessing (boolean, optional): parallelizes four-momenta generations across all available cpu cores. 
    - with_prec (boolean, optional): toggles arbitrary precision; default: float64 precision.
    - loud (boolean, optional): triggers printing of timing statistics.
    - prune_fn (Callable, optional): specifies pruning function to constrain accepted decays.
    - thresh_low (Decimal, optional): specifies lower bound for kinetic energy (eV) of **final** particle configured in Decay initialization.
    - thresh_high (Decimal, optional): specifies upper bound fpr kinetic energy (eV) of **final** particle configured in Decay initialization.
    - reject_stats (boolean, optional): returns statistics for number of decays rejected and accepted during simulation.
    - **smear_and_err_kwargs: arguments for smearing/resolution functions. Optional if set_[type]_kwargs() previously specified.
'''

''' with_prec=False (float64 precision). Returns 1 dictionary, with particle names as keys, and an array of four-momenta as the corresponding value. '''
total_particles =tritium_decay.simulate_decay(parent_name='T', parent_four_momentum=parent_momentum, 
                                                num_batches=10, smearing=True, matrix2_weighting=True, matrix2_upper_estimate_samples=1000, multiprocessing=False, with_prec=False, loud=True, thresh_low=Decimal(18000))

''' with_prec=True (arbitrary precision). Returns 2 dictionaries, with particle names as keys, and an array of four-momenta as the corresponding value.
First dictionary is float64, second dictionary is arbitrary precision. '''
total_particles, total_particles_precise = tritium_decay.simulate_decay(parent_name='T', with_prec=True,...)

Important note: If multiprocessing=True, wrap the simulate_decay call under if __name__ == "__main__":, which is necessary to guard against infinite child processes spawning (primarily for Windows), as documented in multiprocessing library. A warning will be thrown to remind users.

The above precision-based four-momenta dictionaries theoretically scale linearly (multiprocessing yields constant speedup based on number of cpu cores utilized). However, overhead in pickling class instances causes multiprocessing to scale better than linearly up until a threshold based on num_decays.

if __name__ == '__main__':
    total_particles, total_particles_precise = tritium_decay.simulate_decay(parent_name='T', with_prec=True,...)

If reject_stats=True, simulate_decay additionally returns the corresponding accept/reject dictionary.

total_particles, total_particles_precise, reject_dict = tritium_decay.simulate_decay(parent_name='T', with_prec=True, reject_stats=True,...)

The aforementioned four-momenta dictionaries can also be directly accessed after running simulate_decay():

total_particles_precise = tritium_decay.get_total_particles(with_prec=True)

The momenta/energies the particles can also be directly accesssed after running simulate_decay():

'''Float64 numpy array returned if with_prec=False, precise MpMath matrix returned in with_prec=True. '''

# Returns dictionary of energies. Keys are particle names, values are energy arrays. 
energies = tritium_decay.get_energies(particle_names=['e-', 've'], with_prec=False)

# Returns dictionary of momenta vectors. Keys are particle names, values are arrays of momentum vectors [px, py, pz].
momenta_vectors = tritium_decay.get_momenta_vectors(particle_names=['He3'], with_prec=True)

# Returns dictionary of momenta magnitudes. Keys are particle names, values are arrays of momenta magnitudes.
momenta = tritium_decay.get_momenta(particle_names=['e-'], with_prec=True)

Additionally, if smearing=True, one can additionally access the unsmeared four-momenta, via:

# Returns unsmeared four-momenta if smear=True was called in simulate_decay.
pre_smeared = tritium_decay.get_pre_smear_total_particles(with_prec=Trye)

Directly calculate invariants (four-momenta^2):
Can only be run following simulate_decay(). Note that MpMath calculations are relatively slow, so an optimal prec_batches significantly reduces time-complexity. Optimal value for prec_batches is around 0.15 * decays for less than 100K decays, and gradually increases with more decays.

'''
Argument Specifications:
    - loud (boolean, optional): triggers printing of timing statistics.
    - prec_batches (integer, optional): number of batches for speedup; can specify if with_prec=True.
    - prec_batch_fraction (float, optional): batch_size / num_decays; 0 < decimal <= 1; can specify if with_prec=True. Takes precedence over prec_batches.
    - with_prec (boolean, optional): toggles arbitrary precision.
'''

''' with_prec=False (float64 precision). Returns 1 dictionary, with two concatenated particle names as keys, and a list of invariant mass^2 (and cross-terms) as the corresponding value.'''
invariants = tritium_decay.calculate_invariants(with_prec=False, loud=True)

''' with_prec=True (arbitrary precision). Returns 2 dictionaries, with two concatenated particle names as keys, and a list of invariant mass^2 (and cross-terms) as the corresponding value.'''
invariants, invariants_prec = tritium_decay.calculate_invariants(with_prec=True, prec_batches=1500, loud=True)
invariants, invariants_prec = tritium_decay.calculate_invariants(with_prec=True, prec_batch_fraction=0.15)

The aforementioned invariant dictionaries can also be directly accessed after running simulate_decay():

invariants = tritium_decay.get_invariants(with_prec=False)

Following calculate_invariants, kinetic energies can be accessed as well.

# Returns dictionary of kinetic energies. Keys are particle names, values are kinetic energy arrays.
kinetic_energies = tritium_decay.get_kinetic_energies(particle_names=['e-', 'He3'], with_prec=True)

Indirectly calculate an invariant (using conservation of energy):
Can only be run following calculate_invariants().

'''
Argument specifications:
    - particle_name (string): identifier of invariant to calculate.
    - with_prec (boolean, optional): toggles arbitrary precision.
'''

# Returns a list of invariant mass^2 for the particle, with length num_events_to_sample (number of decays)
neutrinoless2_prec = tritium_decay.indirect_calculate_invariant(particle_name='ve', with_prec=True)

Propagate errors:
Linear propagation of errors is currently preferred due for time-complexity reasons. Monte Carlo error propagation is currently very slow. Such time complexity is especially an issue for indirect_calculate_invariant_error(), even with optimal prec_batches (0.15 * mc_samples).

Each 'error' below is actually a (2,) numpy array, with the first element being the -uncertainty^2, and the second element being the +uncertainty^2 (the uncertainty here refers to variance).

'''
Argument specifications:
    - particle_name (string): identifier of particle whose error to analyze.
    - with_prec (boolean, optional): toggles arbitrary precision.
    - linear_prop (boolean, optional). True toggles linear, False toggles monte carlo, error propagation. Default = True.
    - mc_samples (integer, optional): number of Monte Carlo samples per decay; specify if linear_prop=False. Default = 1e4
    - prec_batches (integer, optional): number of batches during monte carlo sampling for speedup; specify if with_prec=True.
    - prec_batch_fraction (float, optional): if with_prec = False: batch_size / num_decays. if with_prec = True: batch_size / mc_samples; 
                                                0 < decimal <= 1; can specify if with_prec=True. Takes precedence over prec_batches.
    - **err_kwargs: Arguments for resolution functions. Optional if set_kwargs() previously specified. 
'''

# Returns an array of errors (variances) for the direct invariant calculation (four momenta^2), with shape (num_decays,2)
direct_errs_electron_monte = tritium_decay.calculate_invariant_error(particle_name='e-', with_prec=True, linear_prop=False, mc_samples=1e2, prec_batch_fraction=0.15)
# Returns an array of errors (variances) for the indirect invariant calculation (conservation of energy), with shape (num_decays, 2).
indirect_errs_neutrino_linear = tritium_decay.indirect_calculate_invariant_error(particle_name='ve', with_prec=True, linear_prop=True)

Histogram Plotting:
Plotting energy/momenta/mass^2.

'''
Argument specifications:
    - particle_names (string list): identifiers of particles to graphs.
    - nrows, ncols (integer, optional): number of rows and columns in output graph. length of particle_names should equal nrows * ncols.
    - bins (integer, optional): number of histogram bins.
    - with_prec (boolean, optional): toggles arbitrary precision. Default is False.
    - pdf (dictionary, optional): superposes graph of pdfs, specifically for plot_energies() and plot_kinetic_energies()
        Valid pdf must be of form:
            { 
                [particle_name] : [list of pdf functions of signature:  float fn(float val)]
            }
    - pdf_bins (integer, optional): if pdf is inputted, can specify number of discrete bins with which to plot the pdfs.
    - analyze_peaks (boolean, optional): prints mean and std of multipeak distributions which are not resolved by graph, primarily for
        plot_histogram_mass2() and plot_histogram_indirect_mass2().
    - bounds (dictionary, optional): allows user to specify x-bounds of graph.
        Valid bounds must be of form ('lower' and 'upper' need not both be specified): 
            {
                particle_name: { 'lower': [float val], 'upper': [float val] }
            }
'''

# Example dictionaries
pdf = {
        'e-': [pdf_fn_electron1, pdf_fn_electron2],
        've': [pdf_fn_neutrino]
      }

bounds = { 
            've': { 'lower': lower_bound_neutrino},
            'He3': { 'upper': upper_bound_proton}
         }

''' Can run after simulate_decay(). '''
tritium_decay.plot_momenta(['He3', 'e-', 've'], nrows=1, ncols=3, bins=50, with_prec=True, bounds=bounds)
tritium_decay.plot_energies(['He3', 'e-', 've'], with_prec=True, pdf=pdf, pdf_bins=1000)
# Preferrable plot momenta vectors only when simulating relatively small number of decays, for time complexity reasons.
tritium_decay.plot_momenta_vectors(['He3', 'e-', 've'], nrows=3, with_prec=False)

''' Can run after calculate_invariants() / indirect_calculate_invariant(). '''
tritium_decay.plot_kinetic_energies(['He3', 'e-', 've'], ncols=3, bins=500, with_prec=True)
# Plot direct invariants (four momenta^2):
tritium_decay.plot_histogram_mass2(particle_names=['ve', 'e-'], nrows=1, ncols=2, bins=10, analyze_peaks=True, with_prec=True)
# Plot indirect invariants (conservation of energy):
tritium_decay.plot_histogram_indirect_mass2(particle_names=['ve'], analyze_peaks=False, with_prec=False)

Saving/Loading Decay instances
Finally, you may save and load Decay instances, such as when transfering from a high performance cluster to one's local. Saving requires a single line; reloading requires 3 steps, as outlined below.

# Saving decay instance
tritium_decay.save_decay_to(filepath)
''' Method load_decay_from in core.py '''
# Loading decay instance -- Step 1
reloaded_decay = core.load_decay_from(filepath)

Note that simply pickling functions is especially finicky, dependent on the environment which the save/load occurred, so the save/load methods above are robust to environments at the expense of not storing functions. Therefore, after loading the decay, the matrix^2, smearing, and resolution functions must be reinitialized. Furthermore, as masses may be callable functions, the specialized function load_particle_masses() must also be called, automatically reconfiguring the decay.

 # Reload particle masses -- Step 2
reloaded_decay.load_particle_masses(neutron, proton, electron, neutrino)
 """Call matrix2, smearing, and resolution setters here -- Step 3"""

Phasespace generation

Generation of four-momenta from phasespace is provided, excluding precision-focused and multiprocessing modifications, by the phasespace library:

Albert Puig, & Jonas Eschle (2019). phasespace: n-body phase space generation in Python. Journal of Open Source Software.

The license for the phasespace GitHub is available in LICENSE_PHASESPACE.
Modified python files based on this library are phasespace_precise.py, kinematics.py, and random_control.py. The file headers explicitly reference Albert Puig, the original code developer.