EICTrack.cc

#undef NDEBUG

#include <fstream>
#include <vector>

#include <ConstField.h>
#include <EventDisplay.h>
#include <FieldManager.h>
#include <KalmanFitStatus.h>
#include <KalmanFitter.h>
#include <KalmanFitterInfo.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 *makeIPMeasurement() {
    TVectorD hitCoords(3);
    hitCoords = 0;
    TMatrixDSym hitCov(3);
    hitCov = 0;
    hitCov[0][0] = 10;
    hitCov[1][1] = 10;
    hitCov[2][2] = 10;
    return new genfit::SpacepointMeasurement(hitCoords, hitCov, 0, 0, nullptr);
}

genfit::AbsMeasurement *makeMeasurement(eic::EnergyDep *entry) {
    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 quickTrackProcess(genfit::KalmanFitter *fitter, genfit::Track *track, int n_passes = 1) {
    for (auto rep : track->getTrackReps()) {
        for (int i = 0; i < n_passes; i++) {
            track->reverseTrack();
            fitter->processTrackPartially(track, rep, 0, -1);
            track->reverseTrack();
            fitter->processTrackPartially(track, rep, 0, -1);
        }
    }
}

double forwardPVal(genfit::Track *track, genfit::AbsTrackRep *rep) {
    double chi2 = 0;
    double ndf = -1. * rep->getDim();
    for (auto point : track->getPointsWithMeasurement()) {
        auto fi = point->getKalmanFitterInfo(rep);
        if (!fi) return 0;
        auto update = fi->getForwardUpdate();
        if (!update) return 0;
        chi2 += update->getChiSquareIncrement();
        ndf += update->getNdf();
    }
    return std::max(0., ROOT::Math::chisquared_cdf_c(chi2, ndf));
}

void trackEvent(proio::Event *event, const int n_validate, genfit::KalmanFitter *fitter) {
    // build kernel from event
    std::map<uint64_t, std::map<uint64_t, double>> kernel;
    bool kernel_from_event = false;
    for (auto id : event->TaggedEntries("Reconstructed")) {
        auto kernel_ud = dynamic_cast<eic::KernelMatrix *>(event->GetEntry(id));
        if (kernel_ud) {
            kernel_from_event = true;
            int flat_index = 0;
            for (int i = 0; i < kernel_ud->observation_size(); i++) {
                uint64_t i_id = kernel_ud->observation(i);
                auto &row = kernel[i_id];
                for (int j = i; j < kernel_ud->observation_size(); j++) {
                    uint64_t j_id = kernel_ud->observation(j);
                    row[j_id] = kernel_ud->kflat(flat_index);
                    flat_index++;
                }
            }
            break;
        }
    }

    // build available observations list and a map of observations to
    // GenFit measurements
    std::vector<uint64_t> available_obs;
    std::map<uint64_t, genfit::AbsMeasurement *> measurements;
    std::map<uint64_t, double> seed_sort_param;
    for (auto id : event->TaggedEntries("Tracker")) {
        if (kernel_from_event && kernel.count(id) == 0) continue;
        auto obs = dynamic_cast<eic::EnergyDep *>(event->GetEntry(id));
        if (obs) {
            auto measurement = makeMeasurement(obs);
            if (measurement) {
                available_obs.push_back(id);
                measurements[id] = measurement;
                double t_sum = 0;
                double weight_sum = 0;
                for (auto pos : obs->pos()) {
                    double weight = pos.weightmod() + 1;
                    t_sum += pos.mean().t() * weight;
                    weight_sum += weight;
                }
                seed_sort_param[id] = t_sum / weight_sum;
            }
        }
    }
    auto seed_sort_fn = [&](const uint64_t a_id, const uint64_t b_id) {
        return seed_sort_param[a_id] > seed_sort_param[b_id];
    };

    // build tracks
    while (available_obs.size() >= n_validate) {
        // build track seed
        std::vector<uint64_t> track_obs;
        std::sort(available_obs.begin(), available_obs.end(), seed_sort_fn);
        // if (seed_sort_param[available_obs.back()] > 1) break; // limit time window for seed hits
        track_obs.push_back(available_obs.back());
        available_obs.pop_back();
        auto current_id = track_obs.back();
        auto measurement = measurements[track_obs[0]];
        TVector3 ip_pos(0, 0, 0);
        auto meas_coords = measurement->getRawHitCoords();
        TVector3 meas_pos(10 * meas_coords[0], 10 * meas_coords[1], 10 * meas_coords[2]);

        auto event_kernel_sort_fn = [&](const uint64_t a_id, const uint64_t b_id) {
            double kernel_a, kernel_b;
            if (kernel[current_id].count(a_id))
                kernel_a = kernel[current_id][a_id];
            else
                kernel_a = kernel[a_id][current_id];
            if (kernel[current_id].count(b_id))
                kernel_b = kernel[current_id][b_id];
            else
                kernel_b = kernel[b_id][current_id];
            return kernel_a < kernel_b;
        };
        auto backup_kernel_sort_fn = [&](const uint64_t a_id, const uint64_t b_id) {
            auto a_coords = measurements[a_id]->getRawHitCoords();
            TVector3 a_pos(a_coords[0], a_coords[1], a_coords[2]);
            auto b_coords = measurements[b_id]->getRawHitCoords();
            TVector3 b_pos(b_coords[0], b_coords[1], b_coords[2]);
            return meas_pos.Dot(a_pos) / a_pos.Mag() < meas_pos.Dot(b_pos) / b_pos.Mag();
        };

        auto track = new genfit::Track;
        track->setStateSeed(ip_pos, meas_pos);
        track->insertMeasurement(makeIPMeasurement());
        track->getPoint(0)->setSortingParameter(0.0);
        track->addTrackRep(new genfit::RKTrackRep(11));
        track->addTrackRep(new genfit::RKTrackRep(13));
        std::map<genfit::TrackPoint *, uint64_t> id_lookup;
        int n_seed = 4;
        for (int i = 0; i < n_seed - 1; i++) {
            if (kernel_from_event)
                std::sort(available_obs.begin(), available_obs.end(), event_kernel_sort_fn);
            else
                std::sort(available_obs.begin(), available_obs.end(), backup_kernel_sort_fn);
            auto id = available_obs.back();
            track_obs.push_back(id);
            available_obs.pop_back();

            current_id = id;
            meas_coords = measurements[id]->getRawHitCoords();
            meas_pos = TVector3(0.01 * meas_coords[0], 0.01 * meas_coords[1], 0.01 * meas_coords[2]);
        }
        for (auto id : track_obs) {
            track->insertMeasurement(measurements[id]->clone());
            auto track_point = track->getPoint(-1);
            track_point->setSortingParameter(seed_sort_param[id]);
            id_lookup[track_point] = id;
        }
        track->sort();
        fitter->setMaxIterations(4);
        try {
            fitter->processTrack(track);
        } catch (genfit::Exception e) {
            std::cerr << e.what() << std::endl;
            continue;
        }
        bool do_extend = false;
        for (auto rep : track->getTrackReps())
            if (track->getFitStatus(rep)->isFitConverged()) {
                do_extend = true;
                break;
            }

        // extend track
        while (do_extend && available_obs.size() > 0) {
            current_id = id_lookup[track->getPoint(-1)];
            if (kernel_from_event)
                std::sort(available_obs.begin(), available_obs.end(), event_kernel_sort_fn);
            else {
                meas_coords = track->getPoint(-1)->getRawMeasurement()->getRawHitCoords();
                meas_pos = TVector3(0.01 * meas_coords[0], 0.01 * meas_coords[1], 0.01 * meas_coords[2]);
                std::sort(available_obs.begin(), available_obs.end(), backup_kernel_sort_fn);
            }

            auto id = available_obs.back();
            double kernel_val;
            if (kernel[current_id].count(id))
                kernel_val = kernel[current_id][id];
            else
                kernel_val = kernel[id][current_id];
            if (kernel_from_event && kernel_val < 0.1) {
                break;
            }

            track_obs.push_back(id);
            available_obs.pop_back();
            auto measurement = measurements[id];
            auto coords = measurement->getRawHitCoords();
            track->insertMeasurement(measurement->clone());
            auto track_point = track->getPoint(-1);
            id_lookup[track_point] = id;
            track_point->setSortingParameter(seed_sort_param[id]);
            bool order_changed = track->sort();

            std::vector<genfit::AbsTrackRep *> all_reps(track->getTrackReps());
            std::vector<genfit::AbsTrackRep *> bad_reps;
            for (auto rep : all_reps) {
                try {
                    if (order_changed)
                        fitter->processTrackWithRep(track, rep);
                    else
                        fitter->processTrackPartially(track, rep, -2, -1);
                } catch (genfit::Exception e) {
                    std::cerr << e.what() << std::endl;
                }
                double pval = forwardPVal(track, rep);
                if (pval > 0.001) continue;
                bad_reps.push_back(rep);
            }
            if (bad_reps.size() == all_reps.size()) {
                for (int i = 1; i < track->getNumPoints(); i++) {
                    if (track_point == track->getPoint(i)) {
                        track->deletePoint(i);
                        break;
                    }
                }
                available_obs.push_back(id);
                track_obs.pop_back();
                break;
            }

            for (auto rep : bad_reps) track->deleteTrackRep(track->getIdForRep(rep));
        }
        track->determineCardinalRep();
        auto rep = track->getCardinalRep();

        // fit track
        if (track_obs.size() >= n_validate) {
            fitter->setMaxIterations(128);
            try {
                fitter->processTrackWithRep(track, rep);
            } catch (genfit::Exception e) {
                std::cerr << e.what() << std::endl;
            }

            if (track->getFitStatus(rep)->isFitConverged()) {
                auto eic_track = new eic::Track;
                auto fieldMgr = genfit::FieldManager::getInstance();
                for (unsigned int id = 0; id < track->getNumPoints() - 1; id++) {
                    try {
                        auto point = track->getPoint(id);
                        auto proio_id = id_lookup[point];
                        if (proio_id > 0) eic_track->add_observation(proio_id);
                        auto info = static_cast<genfit::KalmanFitterInfo *>(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 = track->getTrackLen(rep, id, id + 1);

                        auto segment = eic_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);
                    } catch (genfit::Exception e) {
                        std::cerr << e.what() << std::endl;
                    }
                }
                auto point = track->getPoint(track->getNumPoints() - 1);
                eic_track->add_observation(id_lookup[point]);
                auto entryID = event->AddEntry(eic_track, "Tracker");
                event->TagEntry(entryID, "Reconstructed");
                event->TagEntry(entryID, "Vis");

                track_obs.clear();
                std::cout << "successful fit" << std::endl;
            }
        } else
            std::cout << "fit failure" << std::endl;

        // remove largest outlying hit from track
        double max_chi2_inc = 0;
        int max_chi2_inc_i = 1;
        for (int i = 1; i < track_obs.size() + 1; i++) {
            auto fi = static_cast<genfit::KalmanFitterInfo *>(track->getPoint(i)->getKalmanFitterInfo(rep));
            if (!fi) {
                max_chi2_inc_i = i;
                break;
            }
            auto update = fi->getUpdate(1);
            if (!update) {
                max_chi2_inc_i = i;
                break;
            }
            auto chi2_inc = update->getChiSquareIncrement();
        }
        auto max_chi2_inc_id = id_lookup[track->getPoint(max_chi2_inc_i)];
        track_obs.erase(std::remove(track_obs.begin(), track_obs.end(), max_chi2_inc_id), track_obs.end());

        // return unconsumed hits to available list
        for (auto id : track_obs) available_obs.push_back(id);

        delete track;
    }

    for (auto idMeasPair : measurements) delete idMeasPair.second;
}

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) {
    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::KalmanFitter fitter;
    fitter.setDeltaPval(0.1);
    fitter.setRelChi2Change(1);
    int comp = 0;
    int n_validate = 6;

    int opt;
    while ((opt = getopt(argc, argv, "c:n:h")) != -1) {
        switch (opt) {
            case 'c':
                comp = atoi(optarg);
                break;
            case 'n':
                n_validate = atoi(optarg);
                if (n_validate < 2) {
                    std::cerr << "n must be >= 2" << std::endl;
                    exit(EXIT_FAILURE);
                }
                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);
    }

    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;
    int n_events = 0;
    while (reader->Next(&event)) {
        std::cout << "event " << n_events << std::endl;
        checkMetadata(&event);
        trackEvent(&event, n_validate, &fitter);
        writer->Push(&event);
        n_events++;
        if (n_events == 20) break;
    }

    delete writer;
    delete reader;
}