Unify findAllSectors functions

Pr/PrAlgorithms/src/PrAddUTHitsTool.cpp

@@ -281,38 +281,14 @@ namespace LHCb::Pr {
   template <typename InputType>
   std::array<BoundariesNominal, totUTLayers>
   PrAddUTHitsTool::findSectors( const InputType& inputs, MiniStates& filteredStates, float signedReCur ) const {
-
-    std::array<LHCb::UT::TrackUtils::ExtrapolatedStates, totUTLayers> extrapStatesArray;
+    assert( m_geomcache.has_value() );
 
     if constexpr ( std::is_same_v<InputType, LongTracks> ) {
-
-      for ( auto& layer : extrapStatesArray ) { layer.reserve( inputs.size() ); }
-
       for ( auto const& track : inputs.simd() ) {
         //---Define the tolerance parameters
-        const auto qoverp    = track.qOverP();
-        const auto pos       = track.StatePos( 0 );
-        const auto dir       = track.StateDir( 0 );
-        const auto bendParam = p_utParam * -signedReCur * qoverp;
-        const auto overp     = abs( qoverp );
-        const auto xTol      = p_xTol + p_xTolSlope.value() * overp;
-        const auto yTol      = p_yTolSlope.value() * overp;
-
-        assert( m_geomcache.has_value() );
-
-        for ( auto&& [layerIndex, layer] : LHCb::range::enumerate( m_geomcache->layers ) ) {
-
-          const auto zLayer   = layer.z;
-          const auto yPredLay = pos.y + ( zLayer - pos.z ) * dir.y;
-          const auto xPredLay = pos.x + ( zLayer - pos.z ) * dir.x + bendParam * ( zLayer - p_zUTField.value() );
-
-          auto eStatesArr = extrapStatesArray[layerIndex].template emplace_back<SIMDWrapper::InstructionSet::Best>();
-
-          eStatesArr.template field<LHCb::UT::TrackUtils::ExtStateTag::xLayer>().set( xPredLay );
-          eStatesArr.template field<LHCb::UT::TrackUtils::ExtStateTag::yLayer>().set( yPredLay );
-          eStatesArr.template field<LHCb::UT::TrackUtils::ExtStateTag::xTol>().set( xTol );
-          eStatesArr.template field<LHCb::UT::TrackUtils::ExtStateTag::yTol>().set( yTol );
-        }
+        const auto qoverp = track.qOverP();
+        const auto pos    = track.StatePos( 0 );
+        const auto dir    = track.StateDir( 0 );
 
         auto fState = filteredStates.emplace_back<SIMDWrapper::InstructionSet::Best>();
         fState.field<LHCb::UT::TrackUtils::MiniStateTag::State>().setPosition( pos.x, pos.y, pos.z );
@@ -321,31 +297,6 @@ namespace LHCb::Pr {
         fState.field<LHCb::UT::TrackUtils::MiniStateTag::index>().set( track.indices() );
       }
     } else {
-
-      for ( auto& layer : extrapStatesArray ) { layer.reserve( simd::size ); }
-
-      const auto bendParam = p_utParam * -signedReCur * inputs.qOverP();
-      const auto overp     = abs( inputs.qOverP() );
-      const auto xTol      = p_xTol + p_xTolSlope.value() * overp;
-      const auto yTol      = p_yTolSlope.value() * overp;
-
-      assert( m_geomcache.has_value() );
-
-      for ( auto&& [layerIndex, layer] : LHCb::range::enumerate( m_geomcache->layers ) ) {
-
-        const auto zLayer   = layer.z;
-        const auto yPredLay = inputs.y() + ( zLayer - inputs.z() ) * inputs.ty();
-        const auto xPredLay =
-            inputs.x() + ( zLayer - inputs.z() ) * inputs.tx() + bendParam * ( zLayer - p_zUTField.value() );
-
-        auto eStatesArr = extrapStatesArray[layerIndex].template emplace_back<SIMDWrapper::InstructionSet::Scalar>();
-
-        eStatesArr.template field<LHCb::UT::TrackUtils::ExtStateTag::xLayer>().set( xPredLay );
-        eStatesArr.template field<LHCb::UT::TrackUtils::ExtStateTag::yLayer>().set( yPredLay );
-        eStatesArr.template field<LHCb::UT::TrackUtils::ExtStateTag::xTol>().set( xTol );
-        eStatesArr.template field<LHCb::UT::TrackUtils::ExtStateTag::yTol>().set( yTol );
-      }
-
       auto fState = filteredStates.emplace_back<SIMDWrapper::InstructionSet::Scalar>();
       fState.field<LHCb::UT::TrackUtils::MiniStateTag::State>().setPosition( inputs.x(), inputs.y(), inputs.z() );
       fState.field<LHCb::UT::TrackUtils::MiniStateTag::State>().setDirection( inputs.tx(), inputs.ty() );
@@ -354,10 +305,25 @@ namespace LHCb::Pr {
     }
 
     auto compBoundsArray =
-        LHCb::UTDAQ::findAllSectors<BoundariesNominal, LHCb::UT::TrackUtils::BoundariesNominalTag::types,
-                                    LHCb::UT::TrackUtils::maxNumRowsBoundariesNominal,
-                                    LHCb::UT::TrackUtils::maxNumColsBoundariesNominal>( extrapStatesArray,
-                                                                                        filteredStates, *m_geomcache );
+        LHCb::UTDAQ::findAllSectorsExtrap<BoundariesNominal, LHCb::UT::TrackUtils::BoundariesNominalTag::types,
+                                          LHCb::UT::TrackUtils::maxNumRowsBoundariesNominal,
+                                          LHCb::UT::TrackUtils::maxNumColsBoundariesNominal>(
+            filteredStates, *m_geomcache,
+            [&]( int layerIndex, simd::float_v x, simd::float_v y, simd::float_v z, simd::float_v tx, simd::float_v ty,
+                 simd::float_v qop ) {
+              const auto bendParam = p_utParam * -signedReCur * qop;
+              const auto overp     = abs( qop );
+
+              const simd::float_v xTol = p_xTol + p_xTolSlope.value() * overp;
+              const simd::float_v yTol = p_yTolSlope.value() * overp;
+
+              const simd::float_v zLayer{m_geomcache->layers[layerIndex].z};
+              const simd::float_v dxDy{m_geomcache->layers[layerIndex].dxDy};
+
+              const simd::float_v xLayer = x + ( zLayer - z ) * tx + bendParam * ( zLayer - p_zUTField.value() );
+              const simd::float_v yLayer = y + ( zLayer - z ) * ty;
+              return std::make_tuple( xLayer, yLayer, xTol, yTol );
+            } );
 
     return compBoundsArray;
   }

Pr/PrVeloUT/src/PrVeloUT.cpp

@@ -569,11 +569,35 @@ namespace LHCb::Pr {
 
     LHCb::UT::TrackUtils::MiniStates filteredStates = getStates( inputTracks, outputTracks );
 
+    simd::float_v invMinPt       = 1.0f / m_minPT.value();
+    simd::float_v invMinMomentum = 1.0f / m_minMomentum.value();
+
+    // -- Used for the calculation of the size of the search windows
+    constexpr const std::array<float, totalUTLayers> normFact{0.95f, 1.0f, 1.36f, 1.41f};
+
     auto compBoundsArray = LHCb::UTDAQ::findAllSectorsExtrap<
         LHCb::UT::TrackUtils::BoundariesNominal, LHCb::UT::TrackUtils::BoundariesNominalTag::types,
         LHCb::UT::TrackUtils::maxNumRowsBoundariesNominal, LHCb::UT::TrackUtils::maxNumColsBoundariesNominal>(
-        filteredStates, geometry, m_minPT.value(), m_minMomentum.value(), m_distToMomentum, m_yTol.value(),
-        m_yTolSlope.value(), m_intraLayerDist.value(), m_minLayers );
+        filteredStates, geometry,
+        [&]( int layerIndex, simd::float_v x, simd::float_v y, simd::float_v z, simd::float_v tx, simd::float_v ty,
+             simd::float_v ) {
+          // -- this 0.002 seems a little odd...
+          const simd::float_v theta     = max( 0.002f, sqrt( tx * tx + ty * ty ) );
+          const simd::float_v invMinMom = min( invMinPt * theta, invMinMomentum );
+
+          const simd::float_v xTol =
+              abs( m_distToMomentum * invMinMom * normFact[layerIndex] ) - abs( tx ) * m_intraLayerDist.value();
+          const simd::float_v yTol = m_yTol.value() + m_yTolSlope.value() * xTol;
+
+          const simd::float_v zGeo{geometry.layers[layerIndex].z};
+          const simd::float_v dxDy{geometry.layers[layerIndex].dxDy};
+
+          const simd::float_v yAtZ   = y + ty * ( zGeo - z );
+          const simd::float_v xLayer = x + tx * ( zGeo - z );
+          const simd::float_v yLayer = yAtZ + yTol * dxDy;
+          return std::make_tuple( xLayer, yLayer, xTol, yTol );
+        },
+        m_minLayers );
 
     auto hitsInLayers = LHCb::make_object_array<LHCb::Pr::UT::Mut::Hits, batchSize>( Zipping::generateZipIdentifier(),
                                                                                      LHCb::getMemResource( evtCtx ) );