SegmentLineFitter.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
 
// Compile this file assuming that FP operations may trap.
// Prevents spurious FPEs in the clang build.
#include <CxxUtils/trapping_fp.h>
CXXUTILS_TRAPPING_FP;
 
#include <MuonPatternHelpers/SegmentLineFitter.h>
#include <MuonPatternEvent/SegmentFitterEventData.h>
 
#include <MuonSpacePoint/CalibratedSpacePoint.h>
#include <MuonSpacePoint/SpacePointPerLayerSorter.h>
 
#include <ActsInterop/Logger.h>
#include <ActsInterop/UnitConverters.h>
#include <ActsCalibBase/CalibrationContext.h>
 
 
#include <format>
 
#include <xAODMuonPrepData/sTgcMeasurement.h>
#include <xAODMuonPrepData/MdtDriftCircle.h>
 
 
namespace MuonR4::SegmentFit{
    using namespace Acts;
    using namespace Acts::UnitLiterals;
 
    using Hit_t = SegmentLineFitter::Hit_t;
    using HitVec_t = SegmentLineFitter::HitVec_t;
    using Result_t = SegmentLineFitter::Result_t;
 
    namespace {
        bool isGoodHit(const MuonR4::CalibratedSpacePoint& hit) {
            using enum MuonR4::CalibratedSpacePoint::State;
            return hit.fitState() == Valid;
        }
        bool isPrecisionHit(const MuonR4::CalibratedSpacePoint& hit) {
            using enum xAOD::UncalibMeasType;
            return isGoodHit(hit) && (
                hit.type() == MdtDriftCircleType || hit.type() == MMClusterType ||
                (hit.type() == sTgcStripType && 
                static_cast<const xAOD::sTgcMeasurement*>(hit.spacePoint()->primaryMeasurement())->channelType() ==
                sTgcIdHelper::sTgcChannelTypes::Strip)
            );
        }
        inline unsigned countPrecHits(const HitVec_t& hits) {
            return std::ranges::count_if(hits, [](const Hit_t& hit){
                    return isPrecisionHit(*hit);
            });
        }
        inline unsigned countPhiHits(const HitVec_t& hits) {
            return std::ranges::count_if(hits, [](const Hit_t& hit){
                return isGoodHit(*hit) && hit->measuresPhi();
            });
        }
        inline void removeBeamSpot(HitVec_t& hits){
            hits.erase(std::remove_if(hits.begin(), hits.end(),
                        [](const Hit_t& a){
                            return a->type() == xAOD::UncalibMeasType::Other;
                        }), hits.end());
        }
        HitVec_t copy(const HitVec_t& hits) {
            HitVec_t copied{};
            copied.reserve(hits.size());
            std::ranges::transform(hits, std::back_inserter(copied), 
                                    [](const Hit_t& hit) { 
                                        return std::make_unique<CalibratedSpacePoint>(*hit);
                                    });
            return copied;
        }
        Result_t copy(const Result_t& toCopy) {
            Result_t toRet{};
            using FitPars_t = SegmentLineFitter::FitPars_t;
            static_cast<FitPars_t&>(toRet) = toCopy;
            toRet.measurements = copy(toCopy.measurements);
            return toRet;
        }
    } 
    SegmentLineFitter::Config::RangeArray 
        SegmentLineFitter::Config::defaultRanges() {
            RangeArray rng{};
            constexpr double spatRang = 10._m;
            constexpr double timeTange = 25._ns;
            using enum ParamDefs;
            rng[toUnderlying(y0)] = std::array{-spatRang, spatRang};
            rng[toUnderlying(x0)] = std::array{-spatRang, spatRang};
            rng[toUnderlying(phi)] = std::array{-179._degree, 179._degree};
            rng[toUnderlying(theta)] = std::array{-85._degree,  85._degree};
            rng[toUnderlying(t0)] = std::array{-timeTange, timeTange};
            return rng;
    }
    SegmentLineFitter::SegmentLineFitter(const std::string& name, Config&& config):
        AthMessaging{name},
        m_fitter{config, makeActsAthenaLogger(this, name)},
        m_cfg{config} {
        m_goodHitSel.connect<isGoodHit>();
    }
    Result_t SegmentLineFitter::callLineFit(const Acts::CalibrationContext& cctx,
                                            const Parameters& startPars,
                                            const Amg::Transform3D& localToGlobal,
                                            HitVec_t&& calibHits) const {
        if (m_cfg.doBeamSpot && countPhiHits(calibHits) > 0) {
            const Amg::Transform3D globToLoc{localToGlobal.inverse()};
            Amg::Vector3D beamSpot{globToLoc.translation()};
            Amg::Vector3D beamLine{globToLoc.linear().col(2)};
            SpacePoint::Cov_t covariance{};
            covariance[toUnderlying(AxisDefs::etaCov)] = square(m_cfg.beamSpotRadius);
            covariance[toUnderlying(AxisDefs::phiCov)] = square(m_cfg.beamSpotLength);
            auto beamSpotSP = std::make_unique<CalibratedSpacePoint>(nullptr, std::move(beamSpot));
            beamSpotSP->setBeamDirection(std::move(beamLine));
            beamSpotSP->setCovariance(std::move(covariance));
            ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__<<": Beam spot constraint "
                            <<Amg::toString(beamSpotSP->localPosition())<<", "<<beamSpotSP->covariance());
            calibHits.push_back(std::move(beamSpotSP));
        }
        
        if (msgLvl(MSG::VERBOSE)) {
            const auto [pos, dir] = makeLine(startPars);
            std::stringstream hitStream{};
            for (const Hit_t& hit : calibHits) {
                hitStream<<"       **** "<< (*hit)<<", pull: "
                <<std::sqrt(SeedingAux::chi2Term(pos, dir, *hit)) <<std::endl;
            }
            ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__ <<": Start segment fit with parameters "
                <<toString(startPars) <<", plane location: "<<Amg::toString(localToGlobal)<<std::endl
                <<hitStream.str());
        }
        FitOpts_t fitOpts{};
        Result_t result{};
        fitOpts.calibContext = cctx;
        fitOpts.calibrator = m_cfg.calibrator;
        fitOpts.selector = m_goodHitSel;
        //check the degrees of freedom before try the fit
        const auto dOF = m_fitter.countDoF(calibHits, fitOpts.selector);
        if((dOF.bending + dOF.nonBending) < startPars.size()){
            return result;
        }
        fitOpts.measurements = std::move(calibHits);
        fitOpts.localToGlobal = localToGlobal;
        fitOpts.startParameters = startPars;
        fitOpts.startParameters[toUnderlying(ParamDefs::t0)] = ActsTrk::timeToActs(fitOpts.startParameters[toUnderlying(ParamDefs::t0)]);
        result = m_fitter.fit(std::move(fitOpts));
        result.parameters[toUnderlying(ParamDefs::t0)] = ActsTrk::timeToAthena(result.parameters[toUnderlying(ParamDefs::t0)]);
        centerAlongWire(result);
        return result;
    }
    std::unique_ptr<Segment>
        SegmentLineFitter::fitSegment(const EventContext& ctx,
                                      const SegmentSeed* parent,
                                      const Parameters& startPars,
                                      const Amg::Transform3D& localToGlobal,
                                      HitVec_t&& calibHits) const {
        
        const Acts::CalibrationContext cctx = ActsTrk::getCalibrationContext(ctx);
        if (m_cfg.visionTool) {
            Result_t preFit{};
            preFit.parameters = startPars;
            preFit.measurements = copy(calibHits);
            auto seedCopy = convertToSegment(localToGlobal, parent, std::move(preFit));
            m_cfg.visionTool->visualizeSegment(ctx, *seedCopy, "Prefit");
        }
        Result_t segFit = callLineFit(cctx, startPars, localToGlobal, std::move(calibHits));
        if (m_cfg.visionTool && segFit.converged) {
            auto seedCopy = convertToSegment(localToGlobal, parent, copy(segFit));
            m_cfg.visionTool->visualizeSegment(ctx, *seedCopy, "Intermediate fit"); 
        }
        if (!removeOutliers(cctx, *parent, localToGlobal, segFit)) {
            return nullptr;
        }          
        if (!plugHoles(cctx, *parent, localToGlobal, segFit)) {           
            return nullptr;
        }
        auto finalSeg = convertToSegment(localToGlobal, parent, std::move(segFit));
        if (m_cfg.visionTool) {
            m_cfg.visionTool->visualizeSegment(ctx, *finalSeg, "Final fit");
        }
        return finalSeg;
    }
    std::unique_ptr<Segment> 
        SegmentLineFitter::convertToSegment(const Amg::Transform3D& locToGlob, 
                                            const SegmentSeed* patternSeed,
                                            Result_t&& data) const {
        const auto [locPos, locDir] = makeLine(data.parameters);
        Amg::Vector3D globPos = locToGlob * locPos;
        Amg::Vector3D globDir = locToGlob.linear()* locDir;
 
        std::ranges::sort(data.measurements, [](const Hit_t& a, const Hit_t& b){
            return a->localPosition().z() < b->localPosition().z(); 
        });
        if (msgLvl(MSG::VERBOSE)) {
            std::stringstream sstr{};
            for (const Hit_t& h : data.measurements) {
                sstr<<"   **** "<<(*h)<<", pull: "
                    <<std::sqrt(SeedingAux::chi2Term(locPos, locDir, *h))<<std::endl;
            }
            ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__ <<": Create new segment "
                        <<toString(data.parameters)<<" in "<<patternSeed->msSector()->identString()<<"built from:\n"<<sstr.str());
        }
 
        auto finalSeg = std::make_unique<Segment>(std::move(globPos), std::move(globDir),
                                                  patternSeed, std::move(data.measurements),
                                                  data.chi2, data.nDoF);
        finalSeg->setCallsToConverge(data.nIter);
        finalSeg->setParUncertainties(std::move(data.covariance));
        if (m_fitter.config().fitT0) {
            finalSeg->setSegmentT0(data.parameters[toUnderlying(ParamDefs::t0)]);
        }
        return finalSeg;
    }
 
    bool SegmentLineFitter::removeOutliers(const Acts::CalibrationContext& cctx,
                                           const SegmentSeed& seed,
                                           const Amg::Transform3D& localToGlobal,
                                           Result_t& fitResult) const {
      
 
        if (countPrecHits(fitResult.measurements) < m_cfg.nPrecHitCut || fitResult.nDoF == 0
            || fitResult.nIter > m_fitter.config().maxIter) {
                for(const auto& meas : fitResult.measurements){
                    ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__<<": Measurement from fitresult is" << (*meas));
                }
            ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__ 
                            <<": No degree of freedom available. What shall be removed?!. nDoF: "
                            <<fitResult.nDoF<<", n-meas: "<<countPrecHits(fitResult.measurements));
            return false;
        }
        if (fitResult.converged && fitResult.chi2 / std::max(fitResult.nDoF, 1ul) < m_cfg.outlierRemovalCut) {
            ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__ <<": The segment "<<toString(fitResult.parameters)
                          <<" is already of good quality "<<fitResult.chi2 / std::max(fitResult.nDoF, 1ul)
                          <<". Don't remove outliers");
            return true;
        }
        if (fitResult.nDoF == 0u){
            return false;
        }
        ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__ <<": Segment "
                       <<toString(fitResult.parameters)<<" is of badish quality. Remove worst hit");
        
        if (m_cfg.doBeamSpot) {
            removeBeamSpot(fitResult.measurements);
        }
        const auto [segPos, segDir] = makeLine(fitResult.parameters);
        std::ranges::sort(fitResult.measurements,
                [&segPos, &segDir](const HitVec_t::value_type& a, const HitVec_t::value_type& b){
                    const double chiSqA = isGoodHit(*a) ? SeedingAux::chi2Term(segPos, segDir, *a) : 0.;
                    const double chiSqB = isGoodHit(*b) ? SeedingAux::chi2Term(segPos, segDir, *b) : 0.;
                    return chiSqA < chiSqB;                   
                });
        fitResult.measurements.back()->setFitState(HitState::Outlier);

        Result_t newAttempt = callLineFit(cctx, fitResult.parameters, 
                                          localToGlobal, std::move(fitResult.measurements));
        if (newAttempt.converged) {
            newAttempt.nIter+=fitResult.nIter;
            fitResult = std::move(newAttempt);
             if (m_cfg.visionTool) {
                const EventContext& ctx{*cctx.get<const EventContext*>()};
                auto seedCopy = convertToSegment(localToGlobal, &seed, copy(fitResult));
                m_cfg.visionTool->visualizeSegment(ctx, *seedCopy, "Bad fit recovery");
            }
        } else {
            fitResult.nIter+=newAttempt.nIter;
            fitResult.measurements = std::move(newAttempt.measurements);
        }
        return removeOutliers(cctx, seed, localToGlobal, fitResult);
    }
 
    void SegmentLineFitter::eraseWrongHits(Result_t& candidate) const {
        auto [segPos, segDir] = makeLine(candidate.parameters); 
        cleanStripLayers(segPos, segDir, candidate.measurements);
        candidate.measurements.erase(std::remove_if(candidate.measurements.begin(), 
                                                    candidate.measurements.end(),
            [&](const HitVec_t::value_type& hit){
                if (hit->fitState() == HitState::Valid) {
                    return false;
                } else if (hit->fitState() == HitState::Duplicate) {
                    return true;
                }
                if (hit->type() == xAOD::UncalibMeasType::MdtDriftCircleType) {
                    const double dist = Amg::lineDistance(segPos, segDir, 
                                                            hit->localPosition(), 
                                                            hit->sensorDirection());
                    const auto* dc = static_cast<const xAOD::MdtDriftCircle*>(hit->spacePoint()->primaryMeasurement());
                    return dist >= dc->readoutElement()->innerTubeRadius();
                }
                return false;
            }), candidate.measurements.end());
    }
    inline void SegmentLineFitter::cleanStripLayers(const Amg::Vector3D& linePos,
                                                    const Amg::Vector3D& lineDir,
                                                    HitVec_t& hits) const {
        const SpacePointPerLayerSorter sorter{};
        std::ranges::sort(hits, [&sorter, &linePos, &lineDir](const Hit_t&a ,const Hit_t& b){
            if (a->isStraw() || b->isStraw()) {
                return !a->isStraw();
            }
            if (a->type() == xAOD::UncalibMeasType::Other || 
                b->type() == xAOD::UncalibMeasType::Other) {
                return a->type() != xAOD::UncalibMeasType::Other;
            }
            const unsigned lay_a = sorter.sectorLayerNum(*a->spacePoint());
            const unsigned lay_b = sorter.sectorLayerNum(*b->spacePoint());
            if (lay_a != lay_b) {
                return lay_a < lay_b;
            }
            return SeedingAux::chi2Term(linePos, lineDir, *a) <
                   SeedingAux::chi2Term(linePos, lineDir, *b);
        });
        for (HitVec_t::iterator itr = hits.begin(); itr != hits.end(); ++itr) {
            const Hit_t& hit_a{*itr};
            if (hit_a->isStraw()){
                break;
            }
            if(hit_a->fitState() == HitState::Duplicate || 
               hit_a->type() == xAOD::UncalibMeasType::Other) {
                continue;
            }
            const unsigned lay_a = sorter.sectorLayerNum(*hit_a->spacePoint());
            for (HitVec_t::iterator itr2 = itr + 1; itr2 != hits.end(); ++itr2) {
                const Hit_t& hit_b{*itr2};
                if (hit_b->type() == xAOD::UncalibMeasType::Other) {
                    continue;
                }
                if (lay_a != sorter.sectorLayerNum(*hit_b->spacePoint())) {
                    break;
                } 
                if ( (hit_a->measuresEta() && hit_b->measuresEta()) ||
                     (hit_a->measuresPhi() && hit_b->measuresPhi())) {
                    ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__ <<": Reject "
                        <<m_cfg.idHelperSvc->toString(hit_b->spacePoint()->identify()) <<" in favour of "
                        <<m_cfg.idHelperSvc->toString(hit_a->spacePoint()->identify()));
                    hit_b->setFitState(HitState::Duplicate);
                }
            }
        }
    }
    bool SegmentLineFitter::plugHoles(const Acts::CalibrationContext& cctx,
                                      const SegmentSeed& seed,
                                      const Amg::Transform3D& localToGlobal,
                                      Result_t& toRecover) const {
        ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__ <<": segment "<<toString(toRecover.parameters)
                        <<", chi2: "<<toRecover.chi2 /std::max(toRecover.nDoF, 1ul)
                        <<", nDoF: "<<toRecover.nDoF);
        
 
        std::unordered_set<const SpacePoint*> usedSpacePoints{};
        for (auto& hit : toRecover.measurements) {
            ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__ <<": "<<(*hit)<<" is known");
            usedSpacePoints.insert(hit->spacePoint());
        }
        const EventContext& ctx{*cctx.get<const EventContext*>()};
        
        const double timeOff = toRecover.parameters[toUnderlying(ParamDefs::t0)];
        HitVec_t candidateHits{};
        bool hasCandidate{false};
        const auto [locPos, locDir] = makeLine(toRecover.parameters);

        for (const auto& hit : *seed.parentBucket()){            
            if (usedSpacePoints.count(hit.get())){ 
                continue;
            }
            Hit_t calibHit{};
            double pull{-1.};
            if (hit->isStraw()) {
                using namespace Acts::detail::LineHelper;
                const double dist = signedDistance(locPos, locDir, hit->localPosition(), hit->sensorDirection());
                const auto* dc = static_cast<const xAOD::MdtDriftCircle*>(hit->primaryMeasurement());
                // Check whether the tube is crossed by the hit 
                if (Acts::abs(dist) >= dc->readoutElement()->innerTubeRadius()) {
                    continue;
                }
            } else {
                if (!hit->measuresEta() && 
                    std::abs(hit->sensorDirection().dot(hit->localPosition() - 
                        SeedingAux::extrapolateToPlane(locPos,locDir, *hit))) >
                        std::sqrt(hit->covariance()[toUnderlying(AxisDefs::etaCov)])){
                    continue;
                }
                pull = std::sqrt(SeedingAux::chi2Term(locPos, locDir, *hit));
                if (pull > 1.1 * m_cfg.recoveryPull) {
                    continue;
                }
            }
            calibHit = m_cfg.calibrator->calibrate(ctx, hit.get(), locPos, locDir, timeOff);
            pull = std::sqrt(SeedingAux::chi2Term(locPos, locDir, *calibHit));
            if (pull <= m_cfg.recoveryPull) {
                hasCandidate |= calibHit->fitState() == HitState::Valid;
            } else {
                calibHit->setFitState(HitState::Outlier);
            }
            ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__<<": Candidate hit for recovery "
                    <<seed.msSector()->idHelperSvc()->toString(hit->identify())<<", pull: "<<pull);
            candidateHits.push_back(std::move(calibHit));                
        }
        if (!hasCandidate) {
            ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__<<": No space point candidates for recovery were found");
            toRecover.measurements.insert(toRecover.measurements.end(), 
                                          std::make_move_iterator(candidateHits.begin()),
                                          std::make_move_iterator(candidateHits.end()));
            eraseWrongHits(toRecover);
            return toRecover.nDoF > 0;
        }
 
 
        HitVec_t copied = copy(toRecover.measurements);
        HitVec_t copiedCandidates = copy(candidateHits);
        if (m_cfg.doBeamSpot) {
            removeBeamSpot(copied);
        }
 
        candidateHits.insert(candidateHits.end(), 
                             std::make_move_iterator(copied.begin()), 
                             std::make_move_iterator(copied.end()));
 
        cleanStripLayers(locPos, locDir, candidateHits);
        Result_t recovered = callLineFit(cctx, toRecover.parameters, localToGlobal, 
                                         std::move(candidateHits));
        if (!recovered.converged) {
            return false;
        }
        if (recovered.nDoF <= toRecover.nDoF) {
            for (HitVec_t::value_type& hit : copiedCandidates) {
                hit->setFitState(HitState::Outlier);
                toRecover.measurements.push_back(std::move(hit));
            }
            eraseWrongHits(toRecover);
            return true;
        }
        std::vector<const CalibratedSpacePoint*> stripOutliers{};
        stripOutliers.reserve(toRecover.measurements.size());
        ATH_MSG_VERBOSE(__func__<<"() - "<<__LINE__<<":  Before - chi2: "<<toRecover.chi2
                      <<", nDoF "<<toRecover.nDoF<<" <=> after recovery - chi2: "
                      <<recovered.chi2<<", nDoF: "<<recovered.nDoF);
        
        double redChi2 = recovered.chi2 / std::max(recovered.nDoF, 1ul);
        if (redChi2 < m_cfg.outlierRemovalCut || 
            toRecover.nDoF == 0 || redChi2 < toRecover.chi2 / toRecover.nDoF) {
            ATH_MSG_VERBOSE("Accept segment with recovered "<<(recovered.nDoF  - toRecover.nDoF)<<" hits.");
            recovered.nIter += toRecover.nIter;
            toRecover = std::move(recovered);
            unsigned recovLoop{0u};
            while (++recovLoop <= m_cfg.nRecoveryLoops) {                
                copied = copy(toRecover.measurements);
                if (m_cfg.doBeamSpot) {
                    removeBeamSpot(copied);
                }
                const auto [beforePos, beforeDir] = makeLine(toRecover.parameters);
 
                for (HitVec_t::value_type& copyHit : copied) {
                    if (copyHit->fitState() != HitState::Outlier) {
                        continue;
                    }
                    if (std::sqrt(SeedingAux::chi2Term(beforePos, beforeDir, *copyHit)) < m_cfg.recoveryPull) {
                        copyHit->setFitState(HitState::Valid);
                        stripOutliers.push_back(copyHit.get());
                    } 
                }
                // Nothing to recover
                if (stripOutliers.empty()) {
                    break;
                }
                cleanStripLayers(beforePos, beforeDir, copied);
                // Recovery turned out to be duplicates on the same layer
                if (std::ranges::none_of(stripOutliers,[](const CalibratedSpacePoint* sp){
                        return sp->fitState() == HitState::Valid;
                    })) {
                    break;
                }
                stripOutliers.clear();
                recovered = callLineFit(cctx, toRecover.parameters, localToGlobal, std::move(copied));
                if (!recovered.converged) {
                    break;
                }
                if (recovered.nDoF <= toRecover.nDoF) {
                    break;
                }
                redChi2 = recovered.chi2 / std::max(recovered.nDoF, 1ul);
                if (redChi2 <  m_cfg.outlierRemovalCut || redChi2 < toRecover.chi2 / toRecover.nDoF) {
                    recovered.nIter += toRecover.nIter;
                    toRecover = std::move(recovered);
                } else {
                    break;
                }
            }
        } else{
            for (HitVec_t::value_type& hit : copiedCandidates) {
                hit->setFitState(HitState::Outlier);
                toRecover.measurements.push_back(std::move(hit));
            }
        }
        eraseWrongHits(toRecover);
        return true;
    }
    void SegmentLineFitter::centerAlongWire(Result_t& result) const {
        if (std::ranges::any_of(result.measurements,[](const Hit_t& h){
                return h->measuresPhi();
            })) {
            ATH_MSG_VERBOSE("The segment has phi measurements. No centering needed");
            return;
        }
        double avgX{0.};
        const double nHits = result.measurements.size();
        std::ranges::for_each(result.measurements, [&avgX, nHits](const Hit_t& hit) {
            return avgX += hit->localPosition().x() / nHits;
        });
        result.parameters[toUnderlying(ParamDefs::x0)] = avgX;
    }
         
   
}