File indexing completed on 2025-08-06 08:18:06
0001
0002
0003
0004
0005
0006
0007 #include "TpcSpaceChargeMatrixInversion.h"
0008 #include "TpcSpaceChargeMatrixContainerv2.h"
0009 #include "TpcSpaceChargeReconstructionHelper.h"
0010
0011 #include <frog/FROG.h>
0012 #include <tpc/TpcDistortionCorrectionContainer.h>
0013
0014 #include <TFile.h>
0015 #include <TH2.h>
0016 #include <TH3.h>
0017
0018 #include <Eigen/Core>
0019 #include <Eigen/Dense>
0020
0021 #include <memory>
0022
0023 namespace
0024 {
0025
0026 static constexpr float m_phimin = 0;
0027 static constexpr float m_phimax = 2. * M_PI;
0028
0029
0030
0031 static constexpr float m_rmin = 20;
0032 static constexpr float m_rmax = 78;
0033
0034
0035 static constexpr float m_zmin = -105.5;
0036 static constexpr float m_zmax = 105.5;
0037
0038
0039 template<float (TpcSpaceChargeMatrixContainer::*accessor)(int , int , int ) const, int N>
0040 Eigen::Matrix<float, N, N> get_matrix( const TpcSpaceChargeMatrixContainer* container, int icell )
0041 {
0042 Eigen::Matrix<float, N, N> out;
0043 for( int i = 0; i < N; ++i )
0044 {
0045 for( int j = 0; j < N; ++j )
0046 {
0047 out(i, j) = (container->*accessor)(icell, i, j);
0048 }
0049 }
0050
0051 return out;
0052 }
0053
0054
0055 template<float (TpcSpaceChargeMatrixContainer::*accessor)(int ,int ) const, int N>
0056 Eigen::Matrix<float, N, 1> get_column( const TpcSpaceChargeMatrixContainer* container, int icell )
0057 {
0058 Eigen::Matrix<float, N, 1> out;
0059 for( int i = 0; i < N; ++i )
0060 {
0061 out(i) = (container->*accessor)(icell, i);
0062 }
0063
0064 return out;
0065 }
0066
0067 }
0068
0069
0070 TpcSpaceChargeMatrixInversion::TpcSpaceChargeMatrixInversion(const std::string& name)
0071 : Fun4AllBase(name)
0072 {
0073 }
0074
0075
0076 void TpcSpaceChargeMatrixInversion::load_cm_distortion_corrections(const std::string& filename)
0077 {
0078 std::cout << "TpcSpaceChargeMatrixInversion::load_cm_distortion_corrections - loading " << filename << std::endl;
0079
0080
0081 auto distortion_tfile = TFile::Open(filename.c_str());
0082 if (!distortion_tfile)
0083 {
0084 std::cout << "TpcSpaceChargeMatrixInversion::load_cm_distortion_corrections - cannot open " << filename << std::endl;
0085 exit(1);
0086 }
0087
0088
0089 if (!m_dcc_cm)
0090 {
0091 m_dcc_cm.reset(new TpcDistortionCorrectionContainer);
0092 }
0093
0094 const std::array<const std::string, 2> extension = {{"_negz", "_posz"}};
0095 for (int j = 0; j < 2; ++j)
0096 {
0097 m_dcc_cm->m_hDPint[j] = dynamic_cast<TH1*>(distortion_tfile->Get((std::string("hIntDistortionP") + extension[j]).c_str()));
0098 assert(m_dcc_cm->m_hDPint[j]);
0099 m_dcc_cm->m_hDRint[j] = dynamic_cast<TH1*>(distortion_tfile->Get((std::string("hIntDistortionR") + extension[j]).c_str()));
0100 assert(m_dcc_cm->m_hDRint[j]);
0101 m_dcc_cm->m_hDZint[j] = dynamic_cast<TH1*>(distortion_tfile->Get((std::string("hIntDistortionZ") + extension[j]).c_str()));
0102 assert(m_dcc_cm->m_hDZint[j]);
0103 }
0104
0105
0106 m_dcc_cm->m_dimensions = m_dcc_cm->m_hDPint[0]->GetDimension();
0107 assert(m_dcc_cm->m_dimensions == 2);
0108 }
0109
0110
0111 void TpcSpaceChargeMatrixInversion::load_average_distortion_corrections(const std::string& filename)
0112 {
0113 std::cout << "TpcSpaceChargeMatrixInversion::load_average_distortion_corrections - loading " << filename << std::endl;
0114
0115
0116 auto distortion_tfile = TFile::Open(filename.c_str());
0117 if (!distortion_tfile)
0118 {
0119 std::cout << "TpcSpaceChargeMatrixInversion::load_average_distortion_corrections - cannot open " << filename << std::endl;
0120 exit(1);
0121 }
0122
0123
0124 if (!m_dcc_average)
0125 {
0126 m_dcc_average.reset(new TpcDistortionCorrectionContainer);
0127 }
0128
0129 const std::array<const std::string, 2> extension = {{"_negz", "_posz"}};
0130 for (int j = 0; j < 2; ++j)
0131 {
0132 m_dcc_average->m_hDPint[j] = dynamic_cast<TH1*>(distortion_tfile->Get((std::string("hIntDistortionP") + extension[j]).c_str()));
0133 assert(m_dcc_average->m_hDPint[j]);
0134 m_dcc_average->m_hDRint[j] = dynamic_cast<TH1*>(distortion_tfile->Get((std::string("hIntDistortionR") + extension[j]).c_str()));
0135 assert(m_dcc_average->m_hDRint[j]);
0136 m_dcc_average->m_hDZint[j] = dynamic_cast<TH1*>(distortion_tfile->Get((std::string("hIntDistortionZ") + extension[j]).c_str()));
0137 assert(m_dcc_average->m_hDZint[j]);
0138 }
0139
0140
0141 m_dcc_average->m_dimensions = m_dcc_average->m_hDPint[0]->GetDimension();
0142 assert(m_dcc_average->m_dimensions == 3);
0143 }
0144
0145
0146 bool TpcSpaceChargeMatrixInversion::add_from_file(const std::string& shortfilename, const std::string& objectname)
0147 {
0148
0149 FROG frog;
0150 const auto filename = frog.location(shortfilename);
0151
0152
0153 std::unique_ptr<TFile> inputfile(TFile::Open(filename));
0154 if (!inputfile)
0155 {
0156 std::cout << "TpcSpaceChargeMatrixInversion::add_from_file - could not open file " << filename << std::endl;
0157 return false;
0158 }
0159
0160
0161 std::unique_ptr<TpcSpaceChargeMatrixContainer> source(dynamic_cast<TpcSpaceChargeMatrixContainer*>(inputfile->Get(objectname.c_str())));
0162 if (!source)
0163 {
0164 std::cout << "TpcSpaceChargeMatrixInversion::add_from_file - could not find object name " << objectname << " in file " << filename << std::endl;
0165 return false;
0166 }
0167
0168
0169 return add(*source.get());
0170 }
0171
0172
0173 bool TpcSpaceChargeMatrixInversion::add(const TpcSpaceChargeMatrixContainer& source)
0174 {
0175
0176 if (!m_matrix_container)
0177 {
0178 m_matrix_container.reset(new TpcSpaceChargeMatrixContainerv2);
0179
0180
0181 int phibins = 0;
0182 int rbins = 0;
0183 int zbins = 0;
0184 source.get_grid_dimensions(phibins, rbins, zbins);
0185
0186
0187 m_matrix_container->set_grid_dimensions(phibins, rbins, zbins);
0188 }
0189
0190
0191 return m_matrix_container->add(source);
0192 }
0193
0194
0195 void TpcSpaceChargeMatrixInversion::calculate_distortion_corrections(const InversionMode inversionMode )
0196 {
0197 if (!m_matrix_container)
0198 {
0199 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - no distortion matrices loaded. Aborting" << std::endl;
0200 exit(1);
0201 }
0202
0203
0204 int phibins = 0;
0205 int rbins = 0;
0206 int zbins = 0;
0207 m_matrix_container->get_grid_dimensions(phibins, rbins, zbins);
0208
0209
0210 std::unique_ptr<TH3> hentries(new TH3F("hentries_rec", "hentries_rec", phibins, m_phimin, m_phimax, rbins, m_rmin, m_rmax, zbins, m_zmin, m_zmax));
0211 std::unique_ptr<TH3> hphi(new TH3F("hDistortionP_rec", "hDistortionP_rec", phibins, m_phimin, m_phimax, rbins, m_rmin, m_rmax, zbins, m_zmin, m_zmax));
0212 std::unique_ptr<TH3> hz(new TH3F("hDistortionZ_rec", "hDistortionZ_rec", phibins, m_phimin, m_phimax, rbins, m_rmin, m_rmax, zbins, m_zmin, m_zmax));
0213 std::unique_ptr<TH3> hr(new TH3F("hDistortionR_rec", "hDistortionR_rec", phibins, m_phimin, m_phimax, rbins, m_rmin, m_rmax, zbins, m_zmin, m_zmax));
0214
0215
0216 for (const auto& h : {hentries.get(), hphi.get(), hz.get(), hr.get()})
0217 {
0218 h->GetXaxis()->SetTitle("#phi (rad)");
0219 h->GetYaxis()->SetTitle("r (cm)");
0220 h->GetZaxis()->SetTitle("z (cm)");
0221 }
0222
0223
0224 for (int iphi = 0; iphi < phibins; ++iphi)
0225 {
0226 for (int ir = 0; ir < rbins; ++ir)
0227 {
0228 for (int iz = 0; iz < zbins; ++iz)
0229 {
0230
0231 const auto icell = m_matrix_container->get_cell_index(iphi, ir, iz);
0232
0233
0234 static constexpr int min_cluster_count = 2;
0235 const auto cell_entries = m_matrix_container->get_entries(icell);
0236 if (cell_entries < min_cluster_count)
0237 {
0238 continue;
0239 }
0240
0241 switch( inversionMode )
0242 {
0243 case InversionMode::FullInversion:
0244 {
0245
0246 static constexpr int ncoord = 3;
0247 using matrix_t = Eigen::Matrix<float, ncoord, ncoord>;
0248 using column_t = Eigen::Matrix<float, ncoord, 1>;
0249
0250
0251 matrix_t lhs = get_matrix<&TpcSpaceChargeMatrixContainer::get_lhs,ncoord>(m_matrix_container.get(),icell);
0252 column_t rhs = get_column<&TpcSpaceChargeMatrixContainer::get_rhs,ncoord>(m_matrix_container.get(),icell);
0253
0254 if (Verbosity())
0255 {
0256
0257 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - inverting bin " << iz << ", " << ir << ", " << iphi << std::endl;
0258 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - entries: " << cell_entries << std::endl;
0259 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - lhs: \n"
0260 << lhs << std::endl;
0261 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - rhs: \n"
0262 << rhs << std::endl;
0263 }
0264
0265
0266 const auto cov = lhs.inverse();
0267 auto partialLu = lhs.partialPivLu();
0268 const auto result = partialLu.solve(rhs);
0269
0270
0271 hentries->SetBinContent(iphi + 1, ir + 1, iz + 1, cell_entries);
0272
0273 hphi->SetBinContent(iphi + 1, ir + 1, iz + 1, result(0));
0274 hphi->SetBinError(iphi + 1, ir + 1, iz + 1, std::sqrt(cov(0, 0)));
0275
0276 hz->SetBinContent(iphi + 1, ir + 1, iz + 1, result(1));
0277 hz->SetBinError(iphi + 1, ir + 1, iz + 1, std::sqrt(cov(1, 1)));
0278
0279 hr->SetBinContent(iphi + 1, ir + 1, iz + 1, result(2));
0280 hr->SetBinError(iphi + 1, ir + 1, iz + 1, std::sqrt(cov(2, 2)));
0281
0282 if (Verbosity())
0283 {
0284 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - dphi: " << result(0) << " +/- " << std::sqrt(cov(0, 0)) << std::endl;
0285 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - dz: " << result(1) << " +/- " << std::sqrt(cov(1, 1)) << std::endl;
0286 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - dr: " << result(2) << " +/- " << std::sqrt(cov(2, 2)) << std::endl;
0287 std::cout << std::endl;
0288 }
0289 break;
0290 }
0291
0292 case InversionMode::ReducedInversion_phi:
0293 case InversionMode::ReducedInversion_z:
0294 {
0295
0296 static constexpr int ncoord = 2;
0297 using matrix_t = Eigen::Matrix<float, ncoord, ncoord>;
0298 using column_t = Eigen::Matrix<float, ncoord, 1>;
0299
0300
0301 matrix_t lhs_rphi = get_matrix<&TpcSpaceChargeMatrixContainer::get_lhs_rphi,ncoord>(m_matrix_container.get(),icell);
0302 column_t rhs_rphi = get_column<&TpcSpaceChargeMatrixContainer::get_rhs_rphi,ncoord>(m_matrix_container.get(),icell);
0303 const auto cov_rphi = lhs_rphi.inverse();
0304 auto partialLu_rphi = lhs_rphi.partialPivLu();
0305 const auto result_rphi = partialLu_rphi.solve(rhs_rphi);
0306
0307
0308 matrix_t lhs_z = get_matrix<&TpcSpaceChargeMatrixContainer::get_lhs_z,ncoord>(m_matrix_container.get(),icell);
0309 column_t rhs_z = get_column<&TpcSpaceChargeMatrixContainer::get_rhs_z,ncoord>(m_matrix_container.get(),icell);
0310 const auto cov_z = lhs_z.inverse();
0311 auto partialLu_z = lhs_z.partialPivLu();
0312 const auto result_z = partialLu_z.solve(rhs_z);
0313
0314
0315 hentries->SetBinContent(iphi + 1, ir + 1, iz + 1, cell_entries);
0316
0317 hphi->SetBinContent(iphi + 1, ir + 1, iz + 1, result_rphi(0));
0318 hphi->SetBinError(iphi + 1, ir + 1, iz + 1, std::sqrt(cov_rphi(0, 0)));
0319
0320 hz->SetBinContent(iphi + 1, ir + 1, iz + 1, result_z(0));
0321 hz->SetBinError(iphi + 1, ir + 1, iz + 1, std::sqrt(cov_z(0, 0)));
0322
0323 if( inversionMode == InversionMode::ReducedInversion_phi )
0324 {
0325 hr->SetBinContent(iphi + 1, ir + 1, iz + 1, result_rphi(1));
0326 hr->SetBinError(iphi + 1, ir + 1, iz + 1, std::sqrt(cov_rphi(1, 1)));
0327 } else if( inversionMode == InversionMode::ReducedInversion_z ) {
0328 hr->SetBinContent(iphi + 1, ir + 1, iz + 1, result_z(1));
0329 hr->SetBinError(iphi + 1, ir + 1, iz + 1, std::sqrt(cov_z(1, 1)));
0330 }
0331
0332
0333 if (Verbosity())
0334 {
0335 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - dphi: " << result_rphi(0) << " +/- " << std::sqrt(cov_rphi(0, 0)) << std::endl;
0336 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - dz: " << result_z(0) << " +/- " << std::sqrt(cov_z(0, 0)) << std::endl;
0337 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - dr (rphi): " << result_rphi(1) << " +/- " << std::sqrt(cov_rphi(1, 1)) << std::endl;
0338 std::cout << "TpcSpaceChargeMatrixInversion::calculate_distortion_corrections - dr (z): " << result_z(1) << " +/- " << std::sqrt(cov_z(1, 1)) << std::endl;
0339 std::cout << std::endl;
0340 }
0341 break;
0342 }
0343 }
0344
0345 }
0346 }
0347 }
0348
0349
0350
0351 auto process_histogram = [](TH3* h, const TString& name)
0352 {
0353 const auto& result = TpcSpaceChargeReconstructionHelper::split(h);
0354 std::unique_ptr<TH3> hneg(std::get<0>(result));
0355 std::unique_ptr<TH3> hpos(std::get<1>(result));
0356
0357 return std::make_tuple(
0358 TpcSpaceChargeReconstructionHelper::add_guarding_bins(hneg.get(), name + "_negz"),
0359 TpcSpaceChargeReconstructionHelper::add_guarding_bins(hpos.get(), name + "_posz"));
0360 };
0361
0362
0363 if (!m_dcc_average)
0364 {
0365 m_dcc_average.reset(new TpcDistortionCorrectionContainer);
0366 }
0367
0368 std::tie(m_dcc_average->m_hentries[0], m_dcc_average->m_hentries[1]) = process_histogram(hentries.get(), "hentries");
0369 std::tie(m_dcc_average->m_hDRint[0], m_dcc_average->m_hDRint[1]) = process_histogram(hr.get(), "hIntDistortionR");
0370 std::tie(m_dcc_average->m_hDPint[0], m_dcc_average->m_hDPint[1]) = process_histogram(hphi.get(), "hIntDistortionP");
0371 std::tie(m_dcc_average->m_hDZint[0], m_dcc_average->m_hDZint[1]) = process_histogram(hz.get(), "hIntDistortionZ");
0372 }
0373
0374
0375 void TpcSpaceChargeMatrixInversion::extrapolate_distortion_corrections()
0376 {
0377 if (!m_dcc_average)
0378 {
0379 std::cout << "TpcSpaceChargeMatrixInversion::extrapolate_distortion_corrections - invalid distortion correction container." << std::endl;
0380 return;
0381 }
0382
0383
0384 for (int i = 0; i < 2; ++i)
0385 {
0386 if (!m_dcc_average->m_hDRint[i])
0387 {
0388 std::cout << "TpcSpaceChargeMatrixInversion::extrapolate_distortion_corrections - invalid histograms m_hDRint" << std::endl;
0389 continue;
0390 }
0391
0392
0393
0394 std::unique_ptr<TH3> hmask(static_cast<TH3*>(m_dcc_average->m_hDRint[i]->Clone("hmask")));
0395 TpcSpaceChargeReconstructionHelper::create_tpot_mask(hmask.get());
0396
0397
0398 std::unique_ptr<TH3> hmask_extrap_z(static_cast<TH3*>(hmask->Clone("hmask_extrap_z")));
0399 TpcSpaceChargeReconstructionHelper::extrapolate_z(hmask_extrap_z.get(), hmask.get());
0400
0401
0402
0403
0404
0405 std::unique_ptr<TH3> hmask_extrap_p(static_cast<TH3*>(hmask_extrap_z->Clone("hmask_extrap_p")));
0406 TpcSpaceChargeReconstructionHelper::extrapolate_phi1(hmask_extrap_p.get(), nullptr, hmask_extrap_z.get());
0407
0408
0409
0410
0411
0412 std::unique_ptr<TH3> hmask_extrap_p2(static_cast<TH3*>(hmask_extrap_p->Clone("hmask_extrap_p2")));
0413 TpcSpaceChargeReconstructionHelper::extrapolate_phi2(hmask_extrap_p2.get(), hmask_extrap_p.get());
0414
0415
0416
0417
0418
0419 const uint8_t side = (i == 0) ? TpcSpaceChargeReconstructionHelper::Side_negative : TpcSpaceChargeReconstructionHelper::Side_positive;
0420
0421
0422 auto process_histogram = [&hmask, &hmask_extrap_z, &hmask_extrap_p, &hmask_extrap_p2, side](TH3* h, TH2* h_cm)
0423 {
0424
0425 TpcSpaceChargeReconstructionHelper::extrapolate_z(h, hmask.get());
0426
0427
0428 TpcSpaceChargeReconstructionHelper::extrapolate_phi1(h, h_cm, hmask_extrap_z.get());
0429
0430
0431 TpcSpaceChargeReconstructionHelper::extrapolate_phi2(h, hmask_extrap_p.get());
0432
0433
0434 TpcSpaceChargeReconstructionHelper::extrapolate_z2(h, hmask_extrap_p2.get(), side);
0435 };
0436
0437 process_histogram(static_cast<TH3*>(m_dcc_average->m_hDRint[i]), static_cast<TH2*>(m_dcc_cm->m_hDRint[i]));
0438 process_histogram(static_cast<TH3*>(m_dcc_average->m_hDPint[i]), static_cast<TH2*>(m_dcc_cm->m_hDPint[i]));
0439 process_histogram(static_cast<TH3*>(m_dcc_average->m_hDZint[i]), static_cast<TH2*>(m_dcc_cm->m_hDZint[i]));
0440 }
0441 }
0442
0443
0444 void TpcSpaceChargeMatrixInversion::save_distortion_corrections(const std::string& filename)
0445 {
0446 if (!m_dcc_average)
0447 {
0448 std::cout << "TpcSpaceChargeMatrixInversion::save_distortion_corrections - invalid distortion correction container." << std::endl;
0449 return;
0450 }
0451
0452
0453 std::cout << "TpcSpaceChargeMatrixInversion::save_distortions - writing histograms to " << filename << std::endl;
0454 std::unique_ptr<TFile> outputfile(TFile::Open(filename.c_str(), "RECREATE"));
0455 outputfile->cd();
0456
0457 for (const auto& h_list : {m_dcc_average->m_hentries, m_dcc_average->m_hDRint, m_dcc_average->m_hDPint, m_dcc_average->m_hDZint})
0458 {
0459 for (const auto& h : h_list)
0460 {
0461 if (h)
0462 {
0463 h->Write(h->GetName());
0464 }
0465 }
0466 }
0467
0468
0469 outputfile->Close();
0470 }