EICTrack.cc 21.4 KB
Newer Older
1 2
#undef NDEBUG

David Blyth's avatar
David Blyth committed
3 4 5 6 7 8 9 10
#include <fstream>
#include <vector>

#include <ConstField.h>
#include <EventDisplay.h>
#include <FieldManager.h>
#include <KalmanFitStatus.h>
#include <KalmanFitter.h>
11
#include <KalmanFitterInfo.h>
David Blyth's avatar
David Blyth committed
12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37
#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";
    }
}

David Blyth's avatar
David Blyth committed
38 39 40 41 42 43 44 45 46 47 48
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);
}

David Blyth's avatar
David Blyth committed
49
genfit::AbsMeasurement *makeMeasurement(eic::EnergyDep *entry) {
David Blyth's avatar
David Blyth committed
50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156
    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;
        }
    }
}

David Blyth's avatar
David Blyth committed
157
void quickTrackProcess(genfit::KalmanFitter *fitter, genfit::Track *track, int n_passes = 1) {
David Blyth's avatar
David Blyth committed
158
    for (auto rep : track->getTrackReps()) {
David Blyth's avatar
David Blyth committed
159 160 161 162 163 164
        for (int i = 0; i < n_passes; i++) {
            track->reverseTrack();
            fitter->processTrackPartially(track, rep, 0, -1);
            track->reverseTrack();
            fitter->processTrackPartially(track, rep, 0, -1);
        }
David Blyth's avatar
David Blyth committed
165 166 167 168 169 170 171 172
    }
}

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);
173
        if (!fi) return 0;
David Blyth's avatar
David Blyth committed
174
        auto update = fi->getForwardUpdate();
175
        if (!update) return 0;
David Blyth's avatar
David Blyth committed
176 177 178 179 180 181
        chi2 += update->getChiSquareIncrement();
        ndf += update->getNdf();
    }
    return std::max(0., ROOT::Math::chisquared_cdf_c(chi2, ndf));
}

182
void trackEvent(proio::Event *event, const int n_seed, const int n_validate, genfit::KalmanFitter *fitter) {
David Blyth's avatar
David Blyth committed
183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203
    // 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;
        }
    }

204 205 206 207 208 209
    // 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")) {
David Blyth's avatar
David Blyth committed
210
        if (kernel_from_event && kernel.count(id) == 0) continue;
211 212 213
        auto obs = dynamic_cast<eic::EnergyDep *>(event->GetEntry(id));
        if (obs) {
            auto measurement = makeMeasurement(obs);
David Blyth's avatar
David Blyth committed
214 215 216
            if (measurement) {
                available_obs.push_back(id);
                measurements[id] = measurement;
David Blyth's avatar
David Blyth committed
217 218 219 220 221 222 223 224
                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;
225
            }
226 227 228 229 230 231
        }
    }
    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];
    };

David Blyth's avatar
David Blyth committed
232
    // build tracks
233
    while (available_obs.size() >= n_validate) {
David Blyth's avatar
David Blyth committed
234
        // build track seed
235
        std::vector<uint64_t> track_obs;
David Blyth's avatar
David Blyth committed
236
        std::sort(available_obs.begin(), available_obs.end(), seed_sort_fn);
237
        if (seed_sort_param[available_obs.back()] > 2) break;  // enforce time window for seed hits
David Blyth's avatar
David Blyth committed
238 239
        track_obs.push_back(available_obs.back());
        available_obs.pop_back();
David Blyth's avatar
David Blyth committed
240
        auto current_id = track_obs.back();
David Blyth's avatar
David Blyth committed
241
        auto measurement = measurements[track_obs[0]];
242 243
        TVector3 ip_pos(0, 0, 0);
        auto meas_coords = measurement->getRawHitCoords();
David Blyth's avatar
David Blyth committed
244
        TVector3 meas_pos(10 * meas_coords[0], 10 * meas_coords[1], 10 * meas_coords[2]);
David Blyth's avatar
David Blyth committed
245

246
        // set up sort function for kernel and backup kernel
David Blyth's avatar
David Blyth committed
247 248 249 250 251 252 253 254 255 256 257 258 259
        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) {
David Blyth's avatar
David Blyth committed
260 261 262 263
            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]);
David Blyth's avatar
David Blyth committed
264
            return meas_pos.Dot(a_pos) / a_pos.Mag() < meas_pos.Dot(b_pos) / b_pos.Mag();
David Blyth's avatar
David Blyth committed
265
        };
David Blyth's avatar
David Blyth committed
266 267 268

        auto track = new genfit::Track;
        track->setStateSeed(ip_pos, meas_pos);
David Blyth's avatar
David Blyth committed
269 270
        track->insertMeasurement(makeIPMeasurement());
        track->getPoint(0)->setSortingParameter(0.0);
271
        // add electron hypothesis
David Blyth's avatar
David Blyth committed
272
        track->addTrackRep(new genfit::RKTrackRep(11));
273
        // add muon hypothesis
David Blyth's avatar
David Blyth committed
274
        track->addTrackRep(new genfit::RKTrackRep(13));
275
        std::map<genfit::TrackPoint *, uint64_t> id_lookup;
276
        // add additional seed hits using the kernel to sort by ascending similarity
David Blyth's avatar
David Blyth committed
277
        for (int i = 0; i < n_seed - 1; i++) {
278 279 280 281 282 283
            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);
David Blyth's avatar
David Blyth committed
284
            available_obs.pop_back();
285

286
            // update current hit variables to update the kernel sort functions
287 288 289
            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]);
David Blyth's avatar
David Blyth committed
290
        }
David Blyth's avatar
David Blyth committed
291 292
        for (auto id : track_obs) {
            track->insertMeasurement(measurements[id]->clone());
David Blyth's avatar
David Blyth committed
293 294 295
            auto track_point = track->getPoint(-1);
            track_point->setSortingParameter(seed_sort_param[id]);
            id_lookup[track_point] = id;
David Blyth's avatar
David Blyth committed
296 297
        }
        track->sort();
298
        fitter->setMaxIterations(4);
299 300 301 302 303 304
        try {
            fitter->processTrack(track);
        } catch (genfit::Exception e) {
            std::cerr << e.what() << std::endl;
            continue;
        }
David Blyth's avatar
David Blyth committed
305 306 307 308 309 310
        bool do_extend = false;
        for (auto rep : track->getTrackReps())
            if (track->getFitStatus(rep)->isFitConverged()) {
                do_extend = true;
                break;
            }
David Blyth's avatar
David Blyth committed
311

David Blyth's avatar
David Blyth committed
312 313
        // extend track
        while (do_extend && available_obs.size() > 0) {
314
            // sort available hits relative to last hit in track
315
            current_id = id_lookup[track->getPoint(-1)];
David Blyth's avatar
David Blyth committed
316 317
            if (kernel_from_event)
                std::sort(available_obs.begin(), available_obs.end(), event_kernel_sort_fn);
318 319 320
            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]);
David Blyth's avatar
David Blyth committed
321
                std::sort(available_obs.begin(), available_obs.end(), backup_kernel_sort_fn);
322 323
            }

324 325
            // apply kernel value cutoff to avoid some difficult and
            // unnecessary RK propagations
David Blyth's avatar
David Blyth committed
326
            auto id = available_obs.back();
327 328 329 330 331 332 333 334 335
            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;
            }

David Blyth's avatar
David Blyth committed
336 337 338
            track_obs.push_back(id);
            available_obs.pop_back();
            auto measurement = measurements[id];
David Blyth's avatar
David Blyth committed
339
            auto coords = measurement->getRawHitCoords();
David Blyth's avatar
David Blyth committed
340
            track->insertMeasurement(measurement->clone());
David Blyth's avatar
David Blyth committed
341
            auto track_point = track->getPoint(-1);
David Blyth's avatar
David Blyth committed
342 343
            id_lookup[track_point] = id;
            track_point->setSortingParameter(seed_sort_param[id]);
344
            bool order_changed = track->sort();
David Blyth's avatar
David Blyth committed
345

346
            // fit hypotheses and remove bad hypotheses
David Blyth's avatar
David Blyth committed
347
            std::vector<genfit::AbsTrackRep *> all_reps(track->getTrackReps());
348
            std::vector<genfit::AbsTrackRep *> bad_reps;
David Blyth's avatar
David Blyth committed
349
            for (auto rep : all_reps) {
350
                try {
351 352 353 354
                    if (order_changed)
                        fitter->processTrackWithRep(track, rep);
                    else
                        fitter->processTrackPartially(track, rep, -2, -1);
355 356 357
                } catch (genfit::Exception e) {
                    std::cerr << e.what() << std::endl;
                }
David Blyth's avatar
David Blyth committed
358
                double pval = forwardPVal(track, rep);
359
                if (pval > 0.001) continue;
David Blyth's avatar
David Blyth committed
360 361
                bad_reps.push_back(rep);
            }
362 363 364

            // if all hypotheses are bad, return the latest hit to the list of
            // available hits and break
David Blyth's avatar
David Blyth committed
365
            if (bad_reps.size() == all_reps.size()) {
David Blyth's avatar
David Blyth committed
366
                for (int i = 1; i < track->getNumPoints(); i++) {
David Blyth's avatar
David Blyth committed
367 368 369
                    if (track_point == track->getPoint(i)) {
                        track->deletePoint(i);
                        break;
David Blyth's avatar
David Blyth committed
370
                    }
371
                }
David Blyth's avatar
David Blyth committed
372
                available_obs.push_back(id);
David Blyth's avatar
David Blyth committed
373 374
                track_obs.pop_back();
                break;
375 376
            }

David Blyth's avatar
David Blyth committed
377
            for (auto rep : bad_reps) track->deleteTrackRep(track->getIdForRep(rep));
378
        }
379 380

        // determine best hypothesis
David Blyth's avatar
David Blyth committed
381
        track->determineCardinalRep();
David Blyth's avatar
David Blyth committed
382
        auto rep = track->getCardinalRep();
383

David Blyth's avatar
David Blyth committed
384 385
        // fit track
        if (track_obs.size() >= n_validate) {
386
            fitter->setMaxIterations(128);
387 388 389 390 391
            try {
                fitter->processTrackWithRep(track, rep);
            } catch (genfit::Exception e) {
                std::cerr << e.what() << std::endl;
            }
392

David Blyth's avatar
David Blyth committed
393
            if (track->getFitStatus(rep)->isFitConverged()) {
394
                auto eic_track = new eic::Track;
David Blyth's avatar
David Blyth committed
395 396 397 398
                auto fieldMgr = genfit::FieldManager::getInstance();
                for (unsigned int id = 0; id < track->getNumPoints() - 1; id++) {
                    try {
                        auto point = track->getPoint(id);
399 400
                        auto proio_id = id_lookup[point];
                        if (proio_id > 0) eic_track->add_observation(proio_id);
David Blyth's avatar
David Blyth committed
401
                        auto info = static_cast<genfit::KalmanFitterInfo *>(point->getFitterInfo(rep));
402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426
                        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);
David Blyth's avatar
David Blyth committed
427 428
                    } catch (genfit::Exception e) {
                        std::cerr << e.what() << std::endl;
429 430
                    }
                }
431 432
                auto point = track->getPoint(track->getNumPoints() - 1);
                eic_track->add_observation(id_lookup[point]);
David Blyth's avatar
David Blyth committed
433 434 435
                auto entryID = event->AddEntry(eic_track, "Tracker");
                event->TagEntry(entryID, "Reconstructed");
                event->TagEntry(entryID, "Vis");
436 437

                track_obs.clear();
438
                std::cout << "successful fit" << std::endl;
David Blyth's avatar
David Blyth committed
439
            }
440 441
        } else
            std::cout << "fit failure" << std::endl;
442

David Blyth's avatar
David Blyth committed
443 444 445 446 447 448 449 450 451 452 453 454 455
        // 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;
456
            }
David Blyth's avatar
David Blyth committed
457
            auto chi2_inc = update->getChiSquareIncrement();
458
        }
David Blyth's avatar
David Blyth committed
459 460
        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());
461

462
        // return unconsumed hits to available list
David Blyth's avatar
David Blyth committed
463
        for (auto id : track_obs) available_obs.push_back(id);
464

465 466 467 468 469 470
        delete track;
    }

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

David Blyth's avatar
David Blyth committed
471 472 473 474 475 476 477 478 479
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) {
David Blyth's avatar
David Blyth committed
480 481 482 483 484 485
    class NullBuffer : public std::streambuf {
       public:
        int overflow(int c) { return c; }
    };
    NullBuffer nullBuffer;
    genfit::errorOut.rdbuf(&nullBuffer);
486 487 488 489

    genfit::MaterialEffects::getInstance()->setNoEffects();
    genfit::MaterialEffects::getInstance()->setEnergyLossBrems(false);
    genfit::KalmanFitter fitter;
490
    // set convergence conditions for track fitting
491 492
    fitter.setDeltaPval(0.1);
    fitter.setRelChi2Change(1);
David Blyth's avatar
David Blyth committed
493
    int comp = 0;
494 495
    int n_seed = 4;
    int n_validate = 8;
496

David Blyth's avatar
David Blyth committed
497
    int opt;
498
    while ((opt = getopt(argc, argv, "c:s:n:h")) != -1) {
David Blyth's avatar
David Blyth committed
499 500 501 502
        switch (opt) {
            case 'c':
                comp = atoi(optarg);
                break;
503 504 505 506 507 508 509
            case 's':
                n_seed = atoi(optarg);
                if (n_seed < 1) {
                    std::cerr << "s must be >= 2" << std::endl;
                    exit(EXIT_FAILURE);
                }
                break;
510 511 512 513 514 515 516
            case 'n':
                n_validate = atoi(optarg);
                if (n_validate < 2) {
                    std::cerr << "n must be >= 2" << std::endl;
                    exit(EXIT_FAILURE);
                }
                break;
David Blyth's avatar
David Blyth committed
517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545
            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);
    }

546
    // loop over events
David Blyth's avatar
David Blyth committed
547
    proio::Event event;
548
    int n_events = 0;
David Blyth's avatar
David Blyth committed
549
    while (reader->Next(&event)) {
550
        std::cout << "event " << n_events << std::endl;
551
        // if new field or geometry information is available, apply it
David Blyth's avatar
David Blyth committed
552
        checkMetadata(&event);
553 554 555
        // fit tracks
        trackEvent(&event, n_seed, n_validate, &fitter);
        // write out modified event
David Blyth's avatar
David Blyth committed
556
        writer->Push(&event);
557
        n_events++;
David Blyth's avatar
David Blyth committed
558 559 560 561 562
    }

    delete writer;
    delete reader;
}