.. _program_listing_file_larflow_KeyPoints_PrepAffinityField.cxx:

Program Listing for File PrepAffinityField.cxx
==============================================

|exhale_lsh| :ref:`Return to documentation for file <file_larflow_KeyPoints_PrepAffinityField.cxx>` (``larflow/KeyPoints/PrepAffinityField.cxx``)

.. |exhale_lsh| unicode:: U+021B0 .. UPWARDS ARROW WITH TIP LEFTWARDS

.. code-block:: cpp

   #include "PrepAffinityField.h"
   
   #include "LArUtil/Geometry.h"
   #include "LArUtil/LArProperties.h"
   #include "larflow/Reco/geofuncs.h"
   
   namespace larflow {
   namespace keypoints {
   
     PrepAffinityField::PrepAffinityField()
       : larcv::larcv_base("PrepAffinityField")
     {
       psce = new larutil::SpaceChargeMicroBooNE;
     }
   
     PrepAffinityField::~PrepAffinityField()
     {
       delete psce;
     }
     
     void PrepAffinityField::process( larcv::IOManager& iolcv,
                                      larlite::storage_manager& ioll,
                                      const larflow::prep::PrepMatchTriplets& match_proposals )
     {
   
       larcv::EventImage2D* ev_instance =
         (larcv::EventImage2D*)iolcv.get_data( larcv::kProductImage2D, "instance" );
   
       larlite::event_mctrack* ev_mctrack =
         (larlite::event_mctrack*)ioll.get_data( larlite::data::kMCTrack, "mcreco" );
       larlite::event_mcshower* ev_mcshower =
         (larlite::event_mcshower*)ioll.get_data( larlite::data::kMCShower, "mcreco" );
   
   
       _match_labels_v.clear();
       _match_labels_v.reserve( match_proposals._triplet_v.size() );
   
       std::cout << "Make labels for " << match_proposals._triplet_v.size() << " triplet proposals" << std::endl;
       std::cout << "  plane[" << 0 <<"] pixels: " << match_proposals._sparseimg_vv[0].size() << std::endl;
       std::cout << "  plane[" << 1 <<"] pixels: " << match_proposals._sparseimg_vv[1].size() << std::endl;
       std::cout << "  plane[" << 2 <<"] pixels: " << match_proposals._sparseimg_vv[2].size() << std::endl;    
       
       for ( size_t itriplet=0; itriplet<match_proposals._triplet_v.size(); itriplet++ ) {
         auto const& triplet = match_proposals._triplet_v.at(itriplet);
         int true_triplet = match_proposals._truth_v[itriplet];
         const std::vector<float>& spacepoint = match_proposals._pos_v[itriplet];
         // std::cout << "(" << itriplet << ") [" << triplet[0] << "," << triplet[1] << "," << triplet[2] << "," << triplet[3] << "] "
         //           << " true=" << true_triplet << std::endl;
   
         std::vector<float> label_v;
         float weight = 0.;
         if ( true_triplet==1 ) {
           std::vector< std::vector<int> > pixlist_v(ev_instance->Image2DArray().size());
           for (size_t p=0; p<3; p++) {
             int row = match_proposals._sparseimg_vv[p].at( triplet[p] ).row;
             int col = match_proposals._sparseimg_vv[p].at( triplet[p] ).col;
             pixlist_v[p] = std::vector<int>( { row, col } );
           }
         
           _determine_triplet_labels( pixlist_v,
                                      spacepoint,
                                      ev_instance->Image2DArray(),
                                      *ev_mctrack,
                                      *ev_mcshower,
                                      label_v, weight );
         }
         _match_labels_v.push_back( label_v );
       }
       
     }
   
   
     void PrepAffinityField::_determine_triplet_labels( const std::vector< std::vector<int> >& pixlist_v,
                                                        const std::vector<float>& spacepoint_v,
                                                        const std::vector< larcv::Image2D >& instance_v,
                                                        const larlite::event_mctrack& ev_mctrack_v,
                                                        const larlite::event_mcshower& ev_mcshower_v,
                                                        std::vector<float>& label_v,
                                                        float& weight )
     {
   
       weight = 0.0;
       
       // [0-2]: 3D direction
       // [3-8]: 3 planes x 2D direction
       // [9]: instance id
   
       std::map< int, int > instance_counts;
       for ( size_t p=0; p<pixlist_v.size(); p++ ) {
         //std::cout << "pixel[plane " << p << "] (" << pixlist_v[p][0] << "," << pixlist_v[p][1] << ")" << std::endl;
         int instance_id = (int)instance_v[p].pixel( pixlist_v[p][0], pixlist_v[p][1] );
         //std::cout << "  instanceid=" << instance_id << std::endl;      
         if ( instance_id<=0 )
           continue;
   
         auto it = instance_counts.find(instance_id);
         if ( it==instance_counts.end() ) {
           instance_counts[instance_id] = 0;
           it = instance_counts.find(instance_id);
         }
   
         it->second += 1;
   
       }
   
       int max_id = -1;
       int max_count = 0;
       for ( auto it=instance_counts.begin(); it!=instance_counts.end(); it++ ) {
         if ( it->second>max_count ) {
           max_count = it->second;
           max_id = it->first;
         }
       }
       if ( max_id<=0 || max_count==0 ) {
         return;
       }
   
       //std::cout << "max track-id counts: " << max_id << " counts=" << max_count << std::endl;
   
       int is_track = 0;
       int idx = -1;
   
       for ( int i=0; i<(int)ev_mctrack_v.size(); i++ ) {
         if ( (int)ev_mctrack_v[i].TrackID()==max_id ) {
           is_track = 1;
           idx = i;
           break;
         }
       }
   
       if ( is_track<1 ) {
         for (int i=0; i<(int)ev_mcshower_v.size(); i++ ) {
           if ( (int)ev_mcshower_v[i].TrackID()==max_id  ) {
             is_track = 0;
             idx = i;
           }
         }
       }
   
       if ( idx<0 ) {
         return;
       }
   
       label_v.resize(10,0.0);
       for (int i=0; i<10; i++) label_v[i] = 0.;
       label_v[9] = (float)max_id;
         
       std::vector<double> dspacepoint = { (double)spacepoint_v[0],
                                           (double)spacepoint_v[1],
                                           (double)spacepoint_v[2] };
       
       if ( is_track==1 ) {
         _get_track_direction( ev_mctrack_v[idx], dspacepoint, instance_v, label_v, weight );
       }
       else {
         _get_shower_direction( ev_mcshower_v[idx], dspacepoint, instance_v, label_v, weight );
       }
         
       return;
       
     }
   
     bool PrepAffinityField::_get_track_direction( const larlite::mctrack& track,
                                                   const std::vector<double>& pt,
                                                   const std::vector<larcv::Image2D>& img_v,
                                                   std::vector<float>& label_v,
                                                   float& weight )
     {
   
       int nsteps = track.size();
   
       if ( nsteps==0 ) {
         label_v.clear();
         return false;
       }
       
       std::vector<double> start(3,0);
       std::vector<double> end(3,0);
       std::vector<double> stepdir(3,0);
       double steplen = 0.;
   
       int best_step = -1;
       double best_step_r = -1;
       std::vector<double> best_stepdir(3,0);
       std::vector<double> best_start(3,0);
       std::vector<double> best_end(3,0);
   
       //std::cout << "  track size: " << nsteps << std::endl;
       
       // set start
       for (int i=0; i<3; i++ ) start[i] = track.front().Position()[i];
       std::vector<double> offsets = psce->GetPosOffsets( (double)start[0], (double)start[1], (double)start[2] );
       start[0] -= offsets[0] + 0.7;
       start[1] += offsets[1];
       start[2] += offsets[2];
       
       for (int istep=0; istep+1<nsteps; istep++) {
         // we get the end point, apply the space charge effect
         for (int i=0; i<3; i++ ) end[i] = track[istep+1].Position()[i];
         offsets = psce->GetPosOffsets( (double)end[0], (double)end[1], (double)end[2] );
         end[0] -= offsets[0] + 0.7;
         end[1] += offsets[1];
         end[2] += offsets[2];
   
         // std::cout << "step[" << istep << "] "
         //           << "start=(" << start[0] << "," << start[1] << "," << start[2] << ") -> "
         //           << "end=(" << end[0] << "," << end[1] << "," << end[2] << ")"
         //           << std::endl;
   
         steplen = 0.;
         for (int i=0; i<3; i++ ) {
           stepdir[i] = (end[i]-start[i]);
           steplen += stepdir[i]*stepdir[i];
         }
         if ( steplen>0 ) {
           steplen = sqrt(steplen);
           for (int i=0; i<3; i++ )
             stepdir[i] /= steplen;
         }
   
         if ( steplen>0 ) {
           double projs = larflow::reco::pointRayProjection<double>( start, stepdir, pt );
           if ( projs>0 && projs<steplen ) {
             double r = larflow::reco::pointLineDistance<double>( start, end, pt );
             if ( best_step==-1 || best_step_r>r ) {
               best_step = istep;
               best_step_r = r;
               best_stepdir = stepdir;
               best_start = start;
               best_end = end;
             }
           }
         }
   
         // replace the starting step
         start = end;
       }
   
       if ( best_step==-1 ) {
         // no label generated
         label_v.clear();
         return false;
       }
   
       // copy 3d direction
       for (int i=0; i<3; i++ ) {
         label_v[i] = best_stepdir[i];
       }
   
       // calculate 2d projected direction, by projecting into image, making direction
       // should use plane directions to do this, but projection into plane is the dumb thing to do for now
       float start_tick = 3200 + best_start[0]/larutil::LArProperties::GetME()->DriftVelocity()/0.5;
       float end_tick   = 3200 + best_end[0]/larutil::LArProperties::GetME()->DriftVelocity()/0.5;
       float row_dir = (end_tick-start_tick)/img_v.front().meta().pixel_height();
       for (size_t p=0; p<3; p++ ) {
         float wire_start = larutil::Geometry::GetME()->WireCoordinate( best_start, p );
         float wire_end   = larutil::Geometry::GetME()->WireCoordinate( best_end, p );
         float col_dir = wire_end-wire_start;
         float len2d = 0.;
         len2d = row_dir*row_dir + col_dir*col_dir;
         len2d = sqrt(len2d);
         if ( len2d>0.0 ) {
           label_v[ 3+2*p ]   = row_dir/len2d;
           label_v[ 3+2*p+1 ] = col_dir/len2d;
         }
       }
       
       return true;
     }
   
     bool PrepAffinityField::_get_shower_direction( const larlite::mcshower& shower,
                                                    const std::vector<double>& pt,
                                                    const std::vector<larcv::Image2D>& img_v,
                                                    std::vector<float>& label_v,
                                                    float& weight )
     {
   
       // use det profile direction for 3D label
       float len = 0.;
       for (int i=0; i<3; i++) {
         label_v[i] = shower.DetProfile().Momentum()[i];
         len += label_v[i]*label_v[i];
       }
       len = sqrt(len);
   
       if (len==0) {
         label_v.clear();
         return false;
       }
   
       for (int i=0; i<3; i++) label_v[i] /= len;
   
       // make step
       std::vector<double> start(3,0);
       std::vector<double> end(3,0);
       for (int i=0; i<3; i++) {
         end[i] = start[i] + 1.0*label_v[i];
       }
   
       // sce correct step
       std::vector<double> offset_start = psce->GetPosOffsets( start[0], start[1], start[2] );
       std::vector<double> offset_end   = psce->GetPosOffsets( end[0], end[1], end[2] );
   
       start[0] = start[0] - offset_start[0] + 0.7;
       start[1] = start[1] + offset_start[1];
       start[2] = start[2] + offset_start[2];
   
       end[0] = end[0] - offset_end[0] + 0.7;
       end[1] = end[1] + offset_end[1];
       end[2] = end[2] + offset_end[2];
       
       // calculate the final direction
       len = 0.;
       for (int i=0; i<3; i++ ) {
         label_v[i] = end[i]-start[i];
         len += label_v[i]*label_v[i];
       }
       len = sqrt(len);
   
       if ( len>0 ) {
         for (int i=0; i<3; i++)
           label_v[i] /= len;
       }
       
       // calculate 2d projected direction, by projecting into image, making direction
       // should use plane directions to do this, but projection into plane is the dumb thing to do for now
       float start_tick = 3200 + start[0]/larutil::LArProperties::GetME()->DriftVelocity()/0.5;
       float end_tick   = 3200 + end[0]/larutil::LArProperties::GetME()->DriftVelocity()/0.5;
       float row_dir = (end_tick-start_tick)/img_v.front().meta().pixel_height();
       for (size_t p=0; p<3; p++ ) {
         float wire_start = larutil::Geometry::GetME()->WireCoordinate( start, p );
         float wire_end   = larutil::Geometry::GetME()->WireCoordinate( end, p );
         float col_dir = wire_end-wire_start;
         float len2d = 0.;
         len2d = row_dir*row_dir + col_dir*col_dir;
         len2d = sqrt(len2d);
         if ( len2d>0.0 ) {
           label_v[ 3+2*p ]   = row_dir/len2d;
           label_v[ 3+2*p+1 ] = col_dir/len2d;
         }
       }
   
       return true;
     }
   
     void PrepAffinityField::defineAnaTree()
     {
       _label_tree = new TTree("AffinityFieldTree", "Affinity Field Label Tree");
       _label_tree->Branch( "run",     &_run,    "run/I" );
       _label_tree->Branch( "subrun",  &_subrun, "subrun/I" );
       _label_tree->Branch( "event",   &_event,  "event/I" );
       _label_tree->Branch( "label_v", &_match_labels_v );
     }
   
     void PrepAffinityField::fillAnaTree()
     {
       if (_label_tree) _label_tree->Fill();
     }
   
     void PrepAffinityField::writeAnaTree()
     {
       if ( _label_tree ) _label_tree->Write();
     }
   }
   }