.. _program_listing_file_larflow_KeyPoints_LoaderAffinityField.cxx:

Program Listing for File LoaderAffinityField.cxx
================================================

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

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

.. code-block:: cpp

   #include "LoaderAffinityField.h"
   
   #include <iostream>
   
   namespace larflow {
   namespace keypoints {
   
     bool LoaderAffinityField::_setup_numpy = false;
     
     LoaderAffinityField::LoaderAffinityField( std::vector<std::string>& input_v )
       : tpaf(nullptr)
     {
       input_files.clear();
       input_files = input_v;
       load_tree();
     }
   
     LoaderAffinityField::~LoaderAffinityField()
     {
       if ( tpaf ) delete tpaf;
     }
   
     void LoaderAffinityField::load_tree() {
       std::cout << "[LoaderAffinityField::load_tree()]" << std::endl;
       
       tpaf = new TChain("AffinityFieldTree");
       for (auto const& infile : input_files ) {
         std::cout << "add " << infile << " to AffinityFieldTree chain" << std::endl;
         tpaf->Add(infile.c_str());
       }
       
       _label_v = 0;    
       tpaf->SetBranchAddress( "label_v", &_label_v );
     }
   
     unsigned long LoaderAffinityField::load_entry( int entry )
     {
       unsigned long bytes = tpaf->GetEntry(entry);
       return bytes;
     }
   
     unsigned long LoaderAffinityField::GetEntries()
     {
       return tpaf->GetEntries();
     }
   
     PyObject* LoaderAffinityField::get_match_data( PyObject* triplet_matches_pyobj,
                                                   bool exclude_neg_examples )
     {
   
       _exclude_neg_examples = exclude_neg_examples;
       if ( !_setup_numpy ) {
         import_array1(0);
         _setup_numpy = true;
       }
   
       int index_col = 4;
       
       // cast to numpy array
       PyArrayObject* matches = (PyArrayObject*)triplet_matches_pyobj;
   
       // get the dimensions
       int ndims = PyArray_NDIM( matches );
       npy_intp* shape = PyArray_SHAPE( matches );
       int nfilled = shape[0];
   
       // count npos, nneg examples
       // also make list of indices of positive examples, these are the ones we will evaluate ssnet not
       int npos=0;
       int nneg=0;
       std::vector<int> pos_index_v;
       for (int i=0; i<nfilled; i++) {
         long ispositive = *((long*)PyArray_GETPTR2(matches,i,3));
         if (ispositive==1) {
           npos++;
           pos_index_v.push_back(i);
         }
         else {
           nneg++;
         }
       }
   
       // PARTICLE AFFINITY FIELD GROUND TRUTH
       PyArrayObject* paf_label  = nullptr;
       PyArrayObject* paf_weight = nullptr;
       make_paf_arrays( nfilled, pos_index_v, exclude_neg_examples,
                        matches, paf_label, paf_weight );
       PyObject *paf_label_key     = Py_BuildValue("s", "paf_label" );
       PyObject *paf_weight_key    = Py_BuildValue("s", "paf_weight" );
   
       PyObject *d = PyDict_New();
       PyDict_SetItem(d, paf_label_key,  (PyObject*)paf_label );
       PyDict_SetItem(d, paf_weight_key, (PyObject*)paf_weight ); 
   
       Py_DECREF(paf_label_key);
       Py_DECREF(paf_weight_key);
       
       Py_DECREF(paf_label);
       Py_DECREF(paf_weight);
   
       return d;
     }
   
     int LoaderAffinityField::make_paf_arrays( const int nfilled,
                                              const std::vector<int>& pos_match_index,
                                              const bool exclude_neg_examples,
                                              PyArrayObject* match_array,
                                              PyArrayObject*& paf_label,
                                              PyArrayObject*& paf_weight )
     {
   
       int index_col = 4;
   
       int nd = 2;
       npy_intp dims[] = { (long)pos_match_index.size(), 3 };
       
       if ( !exclude_neg_examples ) {
         dims[0] = (long)nfilled;
       }
       paf_label  = (PyArrayObject*)PyArray_SimpleNew( nd, dims, NPY_FLOAT );
       
       int nd_weight = 1;
       npy_intp dims_weight[] = { dims[0] };
       paf_weight = (PyArrayObject*)PyArray_SimpleNew( nd_weight, dims_weight, NPY_FLOAT ); 
       
       int npos = 0;
       int nneg = 0;
       int nmissing = 0;
   
       std::vector<int> pixtype(dims[0],0);
       for (int i=0; i<(int)dims[0]; i++ ) {
         // sample array index
         int idx = (exclude_neg_examples) ? pos_match_index[i] : (int)i;
         // triplet index
         long index = *((long*)PyArray_GETPTR2(match_array,idx,index_col));
         // ground truth for triplet
         long isgood = *((long*)PyArray_GETPTR2(match_array,idx,3));
   
         const std::vector<float>& label_v = _label_v->at(index);
   
         // if good spacepoint but we don't have a direction label, we zero out the event (missing)
         // if good spacepoint and has direction label, counted as positive example
         // if bad spacepoint point but doesnt matter if have a label, counted as negative example,
         //   correct answer will be zero vector
                         
         if ( isgood==1 && label_v.size()==10 ) {
           // positive examples
           npos++;
           for (int j=0; j<3; j++)
             *((float*)PyArray_GETPTR2(paf_label,i,j)) = label_v[j];
           pixtype[i] = 1;
         }
         else if (isgood==0) {
           // negative examples
           nneg++;
           for (int j=0; j<3; j++)
             *((float*)PyArray_GETPTR2(paf_label,i,j)) = 0.0;
           pixtype[i] = 0;        
         }
         else if (isgood==1 && label_v.size()<10) {
           nmissing++;
           for (int j=0; j<3; j++)
             *((float*)PyArray_GETPTR2(paf_label,i,j)) = 0.0;
           pixtype[i] = 2;        
         }        
       }
       nneg = 0;
   
       // weights for positive and negative examples
       float w_pos = (npos) ? float(npos+nneg)/float(npos) : 0.0;
       float w_neg = (nneg) ? float(npos+nneg)/float(nneg) : 0.0;
       float w_norm = w_pos*npos + w_neg*nneg;
   
       //std::cout << "KPWEIGHT: W(POS)=" << w_pos/w_norm << " W(NEG)=" << w_neg/w_norm << std::endl;
       
       for (int i=0; i<dims[0]; i++ ) {
         if (pixtype[i]==1) {
           *((float*)PyArray_GETPTR1(paf_weight,i)) = w_pos/w_norm;
         }
         else if (pixtype[i]==0) {
           *((float*)PyArray_GETPTR1(paf_weight,i))  = w_neg/w_norm;
         }
         else if (pixtype[i]==2) {
           *((float*)PyArray_GETPTR1(paf_weight,i))  = 0.0;
         }
       }
   
       return 0;
     }
   
     
   }
   }