Q3D-Calibration/include/GaussFit.h
2022-07-06 23:53:09 +08:00

133 lines
3.3 KiB
C++

#pragma once
#ifndef gauss_fit_h
#define gauss_fit_h
#include "GaussNewton.h"
#include "LevenbergMarquardt.h"
#include "clip.h"
#include "utils.h"
#include <TCanvas.h>
#include <TH1F.h>
#include <Eigen/Dense>
#include <iostream>
double Gaussian(double x, double* p) {
return p[0] * std::exp(-(x - p[1]) * (x - p[1]) / (2 * p[2] * p[2]));
}
double* GaussianJacobian(double x, double* p) {
double* resJ = new double[3];
resJ[0] = -Gaussian(x, p) / p[0];
resJ[1] = -(x - p[1]) * Gaussian(x, p) / (p[2] * p[2]);
resJ[2] = -(x - p[1]) * (x - p[1]) * Gaussian(x, p) / (p[2] * p[2] * p[2]);
return resJ;
}
class GaussFit {
public:
GaussFit(){};
~GaussFit(){};
public:
double* parma = new double[3];
std::vector<Eigen::Vector2d> data;
public:
void addData(double x, double y);
double* fit(int type_ = 0);
double RSquare();
int getTotal();
void draw(std::string title = "./Figure.png");
};
void GaussFit::addData(double x, double y) { data.push_back(Eigen::Vector2d(x, y)); }
double* GaussFit::fit(int type_) {
double x, y;
SigmaClip* SC2 = new SigmaClip();
data = SC2->clip(data);
parma[0] = dataMax2D(data);
parma[1] = dataAvg2D(data);
parma[2] = dataStd2D(data);
if (DEBUG > 1)
for (int i = 0; i < data.size(); i++) {
Eigen::Vector2d& point = data.at(i);
x = point(0), y = point(1);
std::cout << x << " " << y << std::endl;
}
if (type_ == 0) {
LevenbergMarquardt LM(3, parma, Gaussian, GaussianJacobian);
LM.data = data;
parma = LM.solve();
} else if (type_ == 1) {
LevenbergMarquardt LM(3, parma, Gaussian, GaussianJacobian, type_ = type_);
LM.data = data;
parma = LM.solve();
} else {
GaussNewton GN(3, parma, Gaussian, GaussianJacobian);
GN.data = data;
parma = GN.solve();
}
if (DEBUG) std::cout << parma[0] << " " << parma[1] << ", " << parma[2] << std::endl;
if (DEBUG) std::cout << RSquare() << std::endl;
if (RSquare() < 0.6)
for (int i = 0; i < 3; i++) parma[i] = 0;
return parma;
}
double GaussFit::RSquare() {
double x, y, mu;
double RSS = 0, TSS = 0;
std::vector<double> yData;
for (int i = 0; i < data.size(); i++) yData.push_back(data.at(i)(1));
mu = dataAvg(yData);
for (int i = 0; i < data.size(); i++) {
Eigen::Vector2d& point = data.at(i);
x = point(0), y = point(1);
RSS += std::pow(y - Gaussian(x, parma), 2);
TSS += std::pow(y - mu, 2);
}
return 1 - RSS / TSS;
}
int GaussFit::getTotal() {
int sum = 0;
for (int i = 0; i < data.size(); i++) {
Eigen::Vector2d& point = data.at(i);
sum += point(1);
}
return sum;
}
void GaussFit::draw(std::string title) {
TCanvas* c1 = new TCanvas("Gauss Fit", "Gauss Fit", 0, 0, 1600, 1200);
TH1F* h1 = new TH1F("", "Raw Data", CHANNEL_NUMBER, 0, CHANNEL_NUMBER);
TH1F* h2 = new TH1F("", "Fit Data", CHANNEL_NUMBER, 0, CHANNEL_NUMBER);
for (int i = 0; i < data.size(); i++) {
Eigen::Vector2d& point = data.at(i);
h1->SetBinContent(point(0), point(1));
}
for (int i = 0; i < CHANNEL_NUMBER; i++) h2->SetBinContent(i, Gaussian(i, parma));
c1->cd(1);
h1->Draw("L");
h2->SetLineColor(kRed);
h2->Draw("SAME");
c1->SaveAs(title.c_str());
}
#endif