Commit 8ae1e154 authored by David Blyth's avatar David Blyth

Initial commit

parents
cmake_minimum_required(VERSION 3.3 FATAL_ERROR)
project(EICTrack VERSION 0.1.0)
set(CMAKE_CXX_STANDARD 11)
# ---------------------------------------------------------------------------
# set default install prefix and build type
if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
set(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT FALSE)
set(CMAKE_INSTALL_PREFIX "/usr/local" CACHE PATH "..." FORCE)
endif()
if(NOT CMAKE_BUILD_TYPE)
set(CMAKE_BUILD_TYPE RelWithDebInfo CACHE STRING "..." FORCE)
endif()
#----------------------------------------------------------------------------
# Find Libraries
find_path(GENFIT_INCLUDE_DIR NAMES ProlateSpacepointMeasurement.h)
if(NOT GENFIT_INCLUDE_DIR)
message(FATAL_ERROR "GenFit headers not found!")
endif()
include_directories(${GENFIT_INCLUDE_DIR})
find_library(GENFIT_LIBRARY NAMES genfit2)
if(NOT GENFIT_LIBRARY)
message(FATAL_ERROR "GenFit was not found!")
endif()
get_filename_component(GENFIT_LIBRARY_DIR ${GENFIT_LIBRARY} DIRECTORY)
link_directories(${GENFIT_LIBRARY_DIR})
find_package(proio 0.10.0 REQUIRED)
include_directories(${proio_INCLUDE_DIR})
find_package(ROOT 6 REQUIRED)
find_package(Eigen3 REQUIRED)
include_directories(${EIGEN3_INCLUDE_DIRS})
#-----------------------------------------------------------------------------------
file(GLOB SRCS
${CMAKE_CURRENT_SOURCE_DIR}/src/EICTrack.cc
${CMAKE_CURRENT_SOURCE_DIR}/src/bfields/*.cc
)
file(GLOB HEADERS
${CMAKE_CURRENT_SOURCE_DIR}/src/EICTrack.h
)
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/src)
add_executable(eictrack ${SRCS} ${HEADERS})
target_link_libraries(eictrack
${GENFIT_LIBRARY}
proio proio.pb
${ROOT_LIBRARIES} ROOT::Geom
${EIGEN3_LIBRARIES}
)
install(TARGETS eictrack
EXPORT ${PROJECT_NAME}Targets
DESTINATION bin
)
#----------------------------------------------------------------------------
# Install and export targets
install(EXPORT ${PROJECT_NAME}Targets
FILE ${PROJECT_NAME}Targets.cmake
DESTINATION lib/${PROJECT_NAME}
)
include(CMakePackageConfigHelpers)
set(TARGETS_INSTALL_PATH lib/${PROJECT_NAME}/${PROJECT_NAME}Targets.cmake)
CONFIGURE_PACKAGE_CONFIG_FILE(
cmake/${PROJECT_NAME}Config.cmake.in
${CMAKE_CURRENT_BINARY_DIR}/${PROJECT_NAME}Config.cmake
INSTALL_DESTINATION lib/${PROJECT_NAME}
PATH_VARS TARGETS_INSTALL_PATH
)
write_basic_package_version_file("${PROJECT_NAME}ConfigVersion.cmake"
VERSION ${${PROJECT_NAME}_VERSION}
COMPATIBILITY SameMajorVersion
)
install(FILES
${CMAKE_CURRENT_BINARY_DIR}/${PROJECT_NAME}Config.cmake
${CMAKE_CURRENT_BINARY_DIR}/${PROJECT_NAME}ConfigVersion.cmake
DESTINATION lib/${PROJECT_NAME}
)
@PACKAGE_INIT@
include("@PACKAGE_TARGETS_INSTALL_PATH@")
check_required_components(${CMAKE_FIND_PACKAGE_NAME})
#include <fstream>
#include <vector>
#include <ConstField.h>
#include <EventDisplay.h>
#include <FieldManager.h>
#include <KalmanFitStatus.h>
#include <KalmanFitter.h>
#include <MaterialEffects.h>
#include <ProlateSpacepointMeasurement.h>
#include <TGeoManager.h>
#include <TGeoMaterialInterface.h>
#include <TVector3.h>
#include <Track.h>
#include <TrackPoint.h>
#include <proio/event.h>
#include <proio/model/eic/eic.pb.h>
#include <proio/reader.h>
#include <proio/writer.h>
#include "EICTrack.h"
using namespace eictrack;
using namespace proio::model;
std::map<std::string, BFieldFactory *> BField::factories;
void printBFields() {
std::cout << "Available magnetic field types:\n";
for (auto type : BField::AvailableTypes()) {
std::cout << "\t" << type << "\n";
}
}
genfit::AbsMeasurement *makeMeasurement(proio::Event *event, uint64_t id, eic::EnergyDep *entry,
double &time) {
bool fromID = false;
time = 0;
int nTimes = 0;
for (auto sourceID : entry->source()) {
auto simHit = dynamic_cast<eic::SimHit *>(event->GetEntry(sourceID));
if (simHit) {
if (simHit->particle() == id) fromID = true;
time += simHit->globalprepos().t();
time += simHit->globalpostpos().t();
nTimes += 2;
}
}
time /= nTimes;
if (!fromID) return NULL;
if (entry->mean() / entry->noise() < 2) return NULL;
TVectorD hitCoords(3);
hitCoords = 0;
TMatrixDSym hitCov(3);
hitCov = 0;
double wTot = 0;
for (auto obsPos : entry->pos()) {
TMatrixDSym hitCovTemp(3);
hitCovTemp = 0;
for (auto randVar : obsPos.noise()) {
hitCovTemp[0][0] += randVar.sigma().x() * randVar.sigma().x();
hitCovTemp[0][1] += randVar.sigma().x() * randVar.sigma().y();
hitCovTemp[0][2] += randVar.sigma().x() * randVar.sigma().z();
hitCovTemp[1][1] += randVar.sigma().y() * randVar.sigma().y();
hitCovTemp[1][2] += randVar.sigma().y() * randVar.sigma().z();
hitCovTemp[2][2] += randVar.sigma().z() * randVar.sigma().z();
}
if (entry->pos().size() == 1) {
hitCov += 0.01 * hitCovTemp;
hitCoords[0] += 0.1 * obsPos.mean().x();
hitCoords[1] += 0.1 * obsPos.mean().y();
hitCoords[2] += 0.1 * obsPos.mean().z();
auto measurement = new genfit::SpacepointMeasurement(hitCoords, hitCov, 0, 0, nullptr);
return measurement;
} else {
hitCovTemp[0][0] += obsPos.mean().x() * obsPos.mean().x();
hitCovTemp[0][1] += obsPos.mean().x() * obsPos.mean().y();
hitCovTemp[0][2] += obsPos.mean().x() * obsPos.mean().z();
hitCovTemp[1][1] += obsPos.mean().y() * obsPos.mean().y();
hitCovTemp[1][2] += obsPos.mean().y() * obsPos.mean().z();
hitCovTemp[2][2] += obsPos.mean().z() * obsPos.mean().z();
double w = obsPos.weightmod() + 1;
hitCov += 0.01 * w * hitCovTemp;
hitCoords[0] += 0.1 * w * obsPos.mean().x();
hitCoords[1] += 0.1 * w * obsPos.mean().y();
hitCoords[2] += 0.1 * w * obsPos.mean().z();
wTot += w;
}
}
hitCov *= 1. / wTot;
hitCoords *= 1. / wTot;
hitCov[0][0] -= hitCoords[0] * hitCoords[0];
hitCov[0][1] -= hitCoords[0] * hitCoords[1];
hitCov[0][2] -= hitCoords[0] * hitCoords[2];
hitCov[1][1] -= hitCoords[1] * hitCoords[1];
hitCov[1][2] -= hitCoords[1] * hitCoords[2];
hitCov[2][2] -= hitCoords[2] * hitCoords[2];
auto measurement = new genfit::SpacepointMeasurement(hitCoords, hitCov, 0, 0, nullptr);
return measurement;
}
void checkMetadata(proio::Event *event) {
static std::shared_ptr<const std::string> lastGeomString(NULL);
if (event->Metadata().count("geometry")) {
auto geomString = event->Metadata().at("geometry");
if (geomString != lastGeomString) {
char dirTemplate[] = "eictrack_XXXXXX";
std::string tempDir = mkdtemp(dirTemplate);
std::string geomFilePath = tempDir + "/geometry.gdml";
std::ofstream geomFile(geomFilePath);
geomFile << *geomString;
geomFile.close();
TGeoManager::Import(geomFilePath.c_str());
genfit::MaterialEffects::getInstance()->setNoEffects(false);
genfit::MaterialEffects::getInstance()->init(new genfit::TGeoMaterialInterface());
unlink(geomFilePath.c_str());
rmdir(tempDir.c_str());
lastGeomString = geomString;
}
}
static std::shared_ptr<const std::string> lastFieldString(NULL);
if (event->Metadata().count("field")) {
auto fieldString = event->Metadata().at("field");
if (fieldString != lastFieldString) {
auto colonLoc = fieldString->find_first_of(':');
auto fieldName = fieldString->substr(0, colonLoc);
std::vector<double> params;
for (int i = colonLoc; i != std::string::npos && i < fieldString->size();
i = fieldString->find_first_of(',', i + 1)) {
double value = atof(fieldString->substr(i + 1).c_str());
params.push_back(value);
}
genfit::AbsBField *field = BField::Create(fieldName, &params[0]);
if (field)
genfit::FieldManager::getInstance()->init(field);
else {
std::cerr << "Invalid magnetic field type: " << fieldName << "\n";
printBFields();
exit(EXIT_FAILURE);
}
lastFieldString = fieldString;
}
}
}
void printUsage() {
std::cerr << "Usage: eictrack [options] <input proio file> <output proio file>\n";
std::cerr << "options:\n";
std::cerr << " -c output compression level: 0 for uncompressed, 1 for LZ4 compression, 2 for GZIP\n"
" compression (default 1))\n";
std::cerr << std::endl;
}
int main(int argc, char **argv) {
int comp = 0;
int opt;
while ((opt = getopt(argc, argv, "c:h")) != -1) {
switch (opt) {
case 'c':
comp = atoi(optarg);
break;
default:
printUsage();
exit(EXIT_FAILURE);
}
}
std::string inputFilename;
std::string outputFilename;
if (optind + 1 < argc) {
inputFilename = argv[optind];
outputFilename = argv[optind + 1];
} else {
printUsage();
exit(EXIT_FAILURE);
}
class NullBuffer : public std::streambuf {
public:
int overflow(int c) { return c; }
};
NullBuffer nullBuffer;
genfit::errorOut.rdbuf(&nullBuffer);
genfit::MaterialEffects::getInstance()->setNoEffects();
genfit::MaterialEffects::getInstance()->setEnergyLossBrems(false);
genfit::AbsKalmanFitter *fitter = new genfit::KalmanFitter();
fitter->setMaxIterations(100);
auto reader = new proio::Reader(inputFilename);
auto writer = new proio::Writer(outputFilename);
switch (comp) {
case 1:
writer->SetCompression(proio::LZ4);
break;
case 2:
writer->SetCompression(proio::GZIP);
break;
default:
writer->SetCompression(proio::UNCOMPRESSED);
}
proio::Event event;
while (reader->Next(&event)) {
checkMetadata(&event);
;
// std::vector<uint64_t> partIDs = event->TaggedEntries("GenStable");
// for (auto partID : partIDs) {
// auto part = dynamic_cast<eic::Particle *>(event->GetEntry(partID));
// if (!part) continue;
//
// std::map<double, genfit::AbsMeasurement *> timeMeasMap;
// auto track = new eic::Track;
//
// std::vector<uint64_t> obsIDs = event->TaggedEntries("Tracker");
// for (auto obsID : obsIDs) {
// auto entry = dynamic_cast<eic::EnergyDep *>(event->GetEntry(obsID));
// if (!entry) continue;
//
// double time;
// auto measurement = makeMeasurement(event, partID, entry, time);
// if (measurement) {
// if (timeMeasMap.count(time) == 0)
// timeMeasMap[time] = measurement;
// else
// delete measurement;
// track->add_observation(obsID);
// }
// }
//
// genfit::AbsTrackRep *rep = new genfit::RKTrackRep(part->pdg());
// TVector3 pos(0, 0, 0);
// TVector3 mom(0, 0, 1);
// if (timeMeasMap.size() > 0) {
// auto iter = timeMeasMap.lower_bound(0);
// const TVectorD &firstHitCoords = iter->second->getRawHitCoords();
// mom[0] = firstHitCoords[0];
// mom[1] = firstHitCoords[1];
// mom[2] = firstHitCoords[2];
// }
// genfit::Track fitTrack(rep, pos, mom);
//
// // assume a vertex plane that goes through the origin
// TVectorD hitCoords(3);
// hitCoords = 0;
// TMatrixDSym hitCov(3);
// hitCov = 0;
// hitCov[0][0] = 10;
// hitCov[1][1] = 10;
// hitCov[2][2] = 10;
// timeMeasMap[0] = new genfit::SpacepointMeasurement(hitCoords, hitCov, 0, 0, nullptr);
//
// for (auto timeMeasPair : timeMeasMap)
// fitTrack.insertPoint(new genfit::TrackPoint(timeMeasPair.second, &fitTrack));
//
// if (timeMeasMap.size() > 4) fitter->processTrack(&fitTrack);
//
// if (fitTrack.getFitStatus()->isFitConverged()) {
// try {
// auto fieldMgr = genfit::FieldManager::getInstance();
// for (unsigned int id = 0; id < fitTrack.getNumPoints() - 1; id++) {
// auto point = fitTrack.getPointWithMeasurement(id);
// auto info = point->getFitterInfo(rep);
// auto state = info->getFittedState();
// TVector3 pos = rep->getPos(state);
// TVector3 mom = rep->getMom(state);
// TVector3 field = fieldMgr->getFieldVal(pos);
// double charge = rep->getCharge(state);
// double chargeMag = abs(charge);
// double time = rep->getTime(state);
// double length = fitTrack.getTrackLen(rep, id, id + 1);
//
// auto segment = track->add_segment();
// auto vertex = segment->mutable_vertex();
// vertex->set_x(pos.x() * 10);
// vertex->set_y(pos.y() * 10);
// vertex->set_z(pos.z() * 10);
// vertex->set_t(time);
// auto poq = segment->mutable_poq();
// poq->set_x(mom.x() / chargeMag);
// poq->set_y(mom.y() / chargeMag);
// poq->set_z(mom.z() / chargeMag);
// auto magfield = segment->mutable_magfield();
// magfield->set_x(field.x() / 10);
// magfield->set_y(field.y() / 10);
// magfield->set_z(field.z() / 10);
// segment->set_chargesign(charge);
// segment->set_length(length * 10);
// }
// auto entryID = event->AddEntry(track, "Tracker");
// event->TagEntry(entryID, "Reconstructed");
// event->TagEntry(entryID, "Vis");
// } catch (const genfit::Exception e) {
// std::cerr << e.what() << std::endl;
// }
// } else
// delete track;
// }
writer->Push(&event);
}
delete writer;
delete reader;
delete fitter;
}
#ifndef CTRACK_GLOBALS_HH
#define CTRACK_GLOBALS_HH
#include <map>
#include <string>
#include <AbsBField.h>
namespace eictrack {
class BFieldFactory {
public:
virtual genfit::AbsBField *Create(double *params) = 0;
};
class BField {
public:
static void RegisterType(const std::string &name, BFieldFactory *factory) { factories[name] = factory; }
static genfit::AbsBField *Create(std::string name, double *params) {
if (factories.count(name))
return factories[name]->Create(params);
else
return NULL;
}
static std::vector<std::string> AvailableTypes() {
std::vector<std::string> types;
for (auto nameFactPair : factories) {
types.push_back(nameFactPair.first);
}
return types;
}
private:
static std::map<std::string, BFieldFactory *> factories;
};
} // namespace eictrack
#define BFIELD_TYPE(klass) \
namespace eictrack { \
class klass##Factory : public BFieldFactory { \
public: \
klass##Factory() { BField::RegisterType(#klass, this); } \
genfit::AbsBField *Create(double *params) { return new klass(params); } \
}; \
static klass##Factory global_##klass##Factory; \
}
#endif // CTRACK_GLOBALS_HH
#include <iostream>
#include "EICTrack.h"
namespace eictrack {
class TOPSiDE : public genfit::AbsBField {
public:
TOPSiDE(double *params) {
_params[0] = params[0] * 10;
_params[1] = params[1] / 10;
_params[2] = params[2] / 10;
_params[3] = params[3] / 10;
_params[4] = params[4] * 10;
_params[5] = params[5] / 10;
_params[6] = params[6] / 10;
_params[7] = params[7] / 10;
}
TVector3 get(const TVector3 &pos) const {
double b[3];
memset(b, 0, 3 * sizeof(double));
get(pos[0], pos[1], pos[2], b[0], b[1], b[2]);
return TVector3(b);
}
void get(const double &posX, const double &posY, const double &posZ, double &Bx, double &By,
double &Bz) const {
double z = posZ;
double r = sqrt(posX * posX + posY * posY);
if (r < _params[3] && z > _params[1] && z < _params[2]) Bz += _params[0];
if (r < _params[7] && z > _params[5] && z < _params[6]) Bx += _params[4];
}
private:
double _params[8];
};
} // namespace eictrack
BFIELD_TYPE(TOPSiDE)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment