NewStructure/waveform_fitting/crystal_ball_fit.cc

0001 #include "waveform_fit_base.h"
0002 #include "crystal_ball_fit.h"
0003
0004 #include <map>
0005 #include <vector>
0006 #include <iostream>
0007
0008 #include <TF1.h>
0009 #include <TGraph.h>
0010 #include <TMath.h>
0011
0012 crystal_ball_fit::crystal_ball_fit() : waveform_fit_base{} {
0013     parameters[0] = 0;  // N
0014     parameters[1] = 0;  // alpha
0015     parameters[2] = 0;  // n
0016     parameters[3] = 0;  // mean
0017     parameters[4] = 0;  // sigma
0018     parameters[5] = 0;  // offset
0019
0020     saturated = false;
0021
0022     crystal_ball_function = new TF1("crystal_ball_function", this, &crystal_ball_fit::crystal_ball, 0, 10, 6);
0023     linear_function = new TF1("constant", "[0]", 0, 10);
0024     crystal_ball_graph = new TGraph();
0025 }
0026
0027 crystal_ball_fit::~crystal_ball_fit() {
0028     delete crystal_ball_function;
0029     delete crystal_ball_graph;
0030 }
0031
0032 void crystal_ball_fit::fit() {
0033     // Fit the waveform with a crystal ball function
0034     // The crystal ball function is defined as:
0035     // f(x; alpha, n, mean, sigma) = N * exp(-0.5 * ((x - mean) / sigma)^2) for x < alpha
0036     //                                N * exp(-0.5 * ((x - mean) / sigma)^2) for x >= alpha
0037     //                                N * (n / abs(alpha)) * (n / abs(alpha)) * exp(-0.5 * alpha * alpha) * (n / abs(alpha) - abs(alpha) - (x - mean) / sigma)^(-n) for x < alpha
0038     //                                N * exp(-0.5 * ((x - mean) / sigma)^2) for x >= alpha
0039     // where N is the normalization factor, alpha is the transition point, n is the power, mean is the mean, and sigma is the standard deviation
0040     // The parameters are stored in the parameters map with the following keys:
0041     // 0: N
0042     // 1: alpha
0043     // 2: n
0044     // 3: mean
0045     // 4: sigma
0046
0047     // Check if the waveform has been set
0048     if (waveform.size() == 0) {
0049         std::cerr << "Waveform has not been set" << std::endl;
0050         return;
0051     }
0052     crystal_ball_function->SetRange(0, waveform.size());
0053
0054     // Fill the graph
0055     for (int i = 0; i < waveform.size(); i++) {
0056         crystal_ball_graph->SetPoint(i, i, waveform[i]);
0057         if (waveform[i] > 1000) {  // Check for saturation.  TODO: Make this a parameter
0058             saturated = true;
0059         }
0060     }
0061
0062     // Try a linear fit
0063     linear_function->SetParameter(0, parameters[5]);
0064     if (parameters.count(15) && parameters.count(25)) {   // If the parameter has a lower and upper limit
0065         linear_function->SetParLimits(0, parameters[15], parameters[25]);
0066     }
0067     crystal_ball_graph->Fit(linear_function, "Q");
0068     if (linear_function->GetChisquare() < 500) {
0069         E = 0;
0070         fit_ndf = linear_function->GetNDF();
0071         fit_chi2 = linear_function->GetChisquare();
0072         stale = false;
0073         return;
0074     }
0075
0076     // Set initial parameters
0077     for (int i = 0; i < 6; i++) {
0078         crystal_ball_function->SetParameter(i, parameters[i]);
0079         if (parameters.count(i + 10) && parameters.count(i + 20)) {   // If the parameter has a lower and upper limit
0080             crystal_ball_function->SetParLimits(i, parameters[i + 10], parameters[i + 20]);
0081         }
0082     }
0083     // Fit the graph
0084     crystal_ball_graph->Fit(crystal_ball_function, "Q");
0085
0086     // Get the results
0087     E = crystal_ball_function->GetParameter(4);
0088     fit_ndf = crystal_ball_function->GetNDF();
0089     fit_chi2 = crystal_ball_function->GetChisquare();
0090
0091     stale = false;
0092 }
0093
0094 int crystal_ball_fit::get_pedestal() {
0095     if (stale) {
0096         return -1;
0097     }
0098     return crystal_ball_function->GetParameter(5);
0099 }
0100
0101 double crystal_ball_fit::crystal_ball(double *inputs, double *parameters) {
0102 // Parameters
0103     // alpha: Where the gaussian transitions to the power law tail - fix?
0104     // n: The exponent of the power law tail - fix?
0105     // x_bar: The mean of the gaussian - free
0106     // sigma: The width of the gaussian - fix ?
0107     // N: The normalization of the gaussian - free
0108     // B baseline - fix?
0109
0110     double x = inputs[0];
0111
0112     double alpha = parameters[0];
0113     double n = parameters[1];
0114     double x_bar = parameters[2];
0115     double sigma = parameters[3];
0116     double N = parameters[4];
0117     double offset = parameters[5];
0118
0119     double A = pow(n / fabs(alpha), n) * exp(-0.5 * alpha * alpha);
0120     double B = n / fabs(alpha) - fabs(alpha);
0121     // std::cout << "A: " << A << std::endl;
0122
0123     // std::cout << "alpha: " << alpha << " n: " << n << " x_bar: " << x_bar << " sigma: " << sigma << " N: " << N << " B: " << B << " A: " << A << std::endl;
0124
0125     double ret_val;
0126     if ((x - x_bar) / sigma < alpha) {
0127         // std::cout << "path a" << std::endl;
0128         ret_val = exp(-0.5 * (x - x_bar) * (x - x_bar) / (sigma * sigma));
0129     } else {
0130         // std::cout << "path b" << std::endl;
0131         ret_val = A * pow(B + (x - x_bar) / sigma, -1 * n);
0132     }
0133     ret_val = N * ret_val + offset;
0134     // std::cout << "x: " << x << " y: " << ret_val << std::endl;
0135     return ret_val;
0136 }