int PHG4HoughTransformTPC::process_event(PHCompositeNode *topNode)
{
  _timer.get()->restart();
  if(_write_reco_tree==true){ _recoevent->tracks.clear();}

  if(verbosity > 0) cout << "PHG4HoughTransformTPC::process_event -- entered" << endl;

  // moving clearing to the beginning of event or we will have
  // return bugs from early exits!
  _clusters_init.clear();
  _clusters.clear();
  _tracks.clear();
  
  //---------------------------------
  // Get Objects off of the Node Tree
  //---------------------------------
  
  GetNodes(topNode);
  
  // Translate into Helix_Hough objects
  //-----------------------------------
  //wrap_clusters_timer.get()->restart();

  for (SvtxClusterMap::Iter iter = _g4clusters->begin();
       iter != _g4clusters->end();
       ++iter) {
    SvtxCluster* cluster = iter->second;

    //cluster->identify();
    
    float phi = atan2(cluster->get_position(1),cluster->get_position(0));
    unsigned int ilayer = _layer_ilayer_map[cluster->get_layer()];
    
    float xy_error=0.;float z_error=0.;
    if (_use_cell_size) {
      xy_error = _smear_xy_layer[ilayer] * _vote_error_scale[ilayer];
      z_error  = _smear_z_layer[ilayer] * _vote_error_scale[ilayer];
      
    }
    else {
      if( cluster->get_phi_size() <= _max_cluster_error*_smear_xy_layer[ilayer] ){xy_error = cluster->get_phi_size() * _vote_error_scale[ilayer];}
      else{xy_error = _max_cluster_error*_smear_xy_layer[ilayer] * _vote_error_scale[ilayer];}
      if(cluster->get_z_size() <= _max_cluster_error*_smear_z_layer[ilayer]){z_error  = cluster->get_z_size() * _vote_error_scale[ilayer];}
      else{z_error  = _max_cluster_error*_smear_z_layer[ilayer] * _vote_error_scale[ilayer];}
    }

    vector<SimpleHit3D>* which_vec = &_clusters;
    if (ilayer<_seed_layers) {which_vec=&_clusters_init;}

    //SimpleHit3D(float xx, float dxx, float yy, float dyy, float zz, float dzz, unsigned int ind, int lyr=-1)
    SimpleHit3D hit3d(cluster->get_x(),fabs(xy_error*sin(phi)),
		      cluster->get_y(),fabs(xy_error*cos(phi)),
		      cluster->get_z(),z_error,
		      cluster->get_id(),ilayer);

    // copy covariance over
    for (int i=0; i<3; ++i) {
      for (int j=i; j<3; ++j) {
	hit3d.set_error(i,j,cluster->get_error(i,j));
      }
    }

    which_vec->push_back(hit3d);
  }

  if (verbosity > 20) {
    cout << "-------------------------------------------------------------------" << endl;
    cout << "PHG4HoughTransformTPC::process_event has the following input clusters:" << endl;

    if (!_clusters_init.empty()) {
      for (unsigned int i = 0; i < _clusters_init.size(); ++i) {
	cout << "n init clusters = "<<_clusters_init.size() << endl;
	_clusters_init[i].print();
      }
    } else {
      for (unsigned int i = 0; i < _clusters.size(); ++i) {
	cout << "n clusters = "<<_clusters.size() << endl;
	_clusters[i].print();
      }
    }
    
    cout << "-------------------------------------------------------------------" << endl;
  }

  cout<<"_clusters_init.size() = "<<_clusters_init.size()<<endl;
  
  //------------------------------------
  // Perform the initial zvertex finding
  //------------------------------------

  if(verbosity > 0) cout << "PHG4HoughTransformTPC::process_event -- initial vertex finding..." << endl;

  // Grab some initial tracks for initial z-vertex finding
  _tracks.clear();

  _vertex.clear();
  _vertex.push_back(0.0); // x guess
  _vertex.push_back(0.0); // y guess
  _vertex.push_back(0.0); // z guess

  if(_use_vertex) {
    
    // find maxtracks tracks
    unsigned int maxtracks = 100;
    // _tracker->setRemoveHits(false);
    _tracker->findHelices(_clusters_init, _req_seed, _max_hits_init, _tracks, maxtracks);
    // _tracker->setRemoveHits(_remove_hits);

    cout<<"found "<<_tracks.size()<<" tracks"<<endl;
    
    if(_tracks.size() == 0){return Fun4AllReturnCodes::EVENT_OK;}
    else if(_tracks.size() == 1)
    {
      _vertex[0] = cos(_tracks[0].phi) * _tracks[0].d;
      _vertex[1] = sin(_tracks[0].phi) * _tracks[0].d;
      _vertex[2] = _tracks[0].z0;
    }
    else
    {
      vector<vector<double> > pTmap;
      for(unsigned int i=0;i<_tracks.size();++i)
      {
        if(_tracks[i].kappa == 0.0){continue;}
        double pT = kappaToPt(_tracks[i].kappa);
        pTmap.push_back(vector<double>());
        pTmap.back().push_back(pT);
        pTmap.back().push_back((double)i);
      }
      sort(pTmap.begin(), pTmap.end());
      vector<SimpleTrack3D> vtxtracks;
      unsigned int maxvtxtracks=100;
      if(_tracks.size() < maxvtxtracks)
      {
        vtxtracks = _tracks;
      }
      else
      {
        for(unsigned int i=0;i<maxvtxtracks;++i)
        {
          vtxtracks.push_back(_tracks[ (int)(pTmap[pTmap.size()-1-i][1]) ]);
        }
      }
      
      
      vector<double> zvertices(3,0.);
      vector<float> temp_vertex(3,0.);
      vector<unsigned int> vtracks(3,0);
      for(unsigned int iter = 0;iter < 3; ++iter)
      {
        temp_vertex[2] = 0.;
        
        TH1D z0_hist("z0_hist","z0_hist", 20, -10., 10.);
        for(unsigned int i=0;i<vtxtracks.size();++i)
        {
          z0_hist.Fill(vtxtracks[i].z0);
        }
        temp_vertex[2] = z0_hist.GetBinCenter( z0_hist.GetMaximumBin() );
        
        _vertexFinder.findVertex(vtxtracks, temp_vertex, 3., true);
        _vertexFinder.findVertex(vtxtracks, temp_vertex, 0.1, true);
        _vertexFinder.findVertex(vtxtracks, temp_vertex, 0.02, false);
        
        
        vector<SimpleTrack3D> ttracks;
        for(unsigned int t=0;t<vtxtracks.size();++t)
        {
          if( fabs(vtxtracks[t].z0 - temp_vertex[2]) < 0.1 ){vtracks[iter] += 1;}
          else{ttracks.push_back(vtxtracks[t]);}
        }
        vtxtracks = ttracks;
        zvertices[iter] = temp_vertex[2];
      }
      _vertex[2] = zvertices[0];
      unsigned int zbest = 0;
      for(unsigned int iter = 1;iter < 3; ++iter)
      {
        if(vtracks[iter] > vtracks[zbest])
        {
          _vertex[2] = zvertices[iter];
          zbest = iter;
        }
      }
    }
    
    
    
    
    
    
    if(verbosity > 0) cout << "PHG4HoughTransformTPC::process_event -- found initial vertex : " << _vertex[0] << " " << _vertex[1] << " " << _vertex[2] << endl;
    
    _tracks.clear();
    
    // shift the vertex to the origin
    for(unsigned int ht=0;ht<_clusters_init.size();++ht)
    {
      _clusters_init[ht].x -= _vertex[0];
      _clusters_init[ht].y -= _vertex[1];
      _clusters_init[ht].z -= _vertex[2];
    }
    for(unsigned int ht=0;ht<_clusters.size();++ht)
    {
      _clusters[ht].x -= _vertex[0];
      _clusters[ht].y -= _vertex[1];
      _clusters[ht].z -= _vertex[2];
    }
    
    

    
  }  // if(_use_vertex)
  
  //----------------------------------
  // Preform the track finding
  //----------------------------------
  _tracker->clear();
  _tracks.clear();
  _timer_initial_hough.get()->restart();
  _tracker->findHelices(_clusters_init, _min_hits_init, _max_hits_init, _tracks);

  _timer_initial_hough.get()->stop();
  
  

  if(verbosity > 0)
  {
    cout << "PHG4HoughTransformTPC::process_event -- full track finding pass found: " << _tracks.size() << " tracks" << endl;
  }    
   
  //----------------------------
  // Re-center event on detector
  //----------------------------

  if(verbosity > 0) cout << "PHG4HoughTransformTPC::process_event -- recentering event on detector..." << endl;
  vector<double> chi_squareds;
  for(unsigned int tt=0;tt<_tracks.size();tt++)
  {
    // move the hits in the track back to their original position                
    for(unsigned int hh=0;hh<_tracks[tt].hits.size();hh++)
    {
      _tracks[tt].hits[hh].x = _tracks[tt].hits[hh].x + _vertex[0];
      _tracks[tt].hits[hh].y = _tracks[tt].hits[hh].y + _vertex[1];
      _tracks[tt].hits[hh].z = _tracks[tt].hits[hh].z + _vertex[2];
      // _tracks[tt].z0 += _vertex[2];
    }
    chi_squareds.push_back(_tracker->getKalmanStates()[tt].chi2);}

  if(verbosity > 0)
  {
    cout << "PHG4HoughTransformTPC::process_event -- final track count: " << _tracks.size() << endl;
  }

  //---------------------------
  // Final vertex determination
  //---------------------------
  
  // final best guess of the primary vertex position here...
  if(verbosity > 0)
  {
    cout<< "PHG4HoughTransformTPC::process_event -- calculating final vertex" << endl;
  }
  
  // sort the tracks by pT
  vector<vector<double> > pTmap;
  for(unsigned int i=0;i<_tracks.size();++i)
  {
    double pT = kappaToPt(_tracks[i].kappa);
    pTmap.push_back(vector<double>());
    pTmap.back().push_back(pT);
    pTmap.back().push_back((double)i);
  }
  sort(pTmap.begin(), pTmap.end());
  vector<SimpleTrack3D> vtxtracks;
  vector<Matrix<float,5,5> > vtxcovariances;
  unsigned int maxvtxtracks=100;
  if(_tracks.size() < maxvtxtracks){vtxtracks = _tracks;}
  else
  {
    for(unsigned int i=0;i<maxvtxtracks;++i)
    {
      vtxtracks.push_back(_tracks[ (int)(pTmap[pTmap.size()-1-i][1]) ]);
      vtxcovariances.push_back( (_tracker->getKalmanStates())[ (int)(pTmap[pTmap.size()-1-i][1]) ].C );
    }
  }
  
  double vx = _vertex[0];
  double vy = _vertex[1];
  double vz = _vertex[2];
  
  _vertex[0] = 0.;
  _vertex[1] = 0.;
  _vertex[2] = 0.;
  
  _vertexFinder.findVertex(vtxtracks, vtxcovariances, _vertex, 0.3, false);
  _vertexFinder.findVertex(vtxtracks, vtxcovariances, _vertex, 0.1, false);
  _vertexFinder.findVertex(vtxtracks, vtxcovariances, _vertex, 0.02, false);
  _vertexFinder.findVertex(vtxtracks, vtxcovariances, _vertex, 0.005, false);
  
  _vertex[0] += vx;
  _vertex[1] += vy;
  _vertex[2] += vz;
  
  if(verbosity > 0)
  {
    cout << "PHG4HoughTransformTPC::process_event -- final vertex: " << _vertex[0] << " " << _vertex[1] << " " << _vertex[2] << endl;
  }

  //--------------------------------
  // Translate back into PHG4 objects
  //--------------------------------

  if(verbosity > 0)
  {
    cout << "PHG4HoughTransformTPC::process_event -- producing PHG4Track objects..." << endl;
  }

  SvtxVertex_v1 vertex;
  vertex.set_t0(0.0);
  for (int i=0;i<3;++i) vertex.set_position(i,_vertex[i]);
  vertex.set_chisq(0.0);
  vertex.set_ndof(0); 
  vertex.set_error(0,0,0.0);
  vertex.set_error(0,1,0.0);
  vertex.set_error(0,2,0.0);
  vertex.set_error(1,0,0.0);
  vertex.set_error(1,1,0.0);
  vertex.set_error(1,2,0.0);
  vertex.set_error(2,0,0.0);
  vertex.set_error(2,1,0.0);
  vertex.set_error(2,2,0.0);
  
  // copy out the reconstructed vertex position
  //_g4tracks->setVertex(_vertex[0],_vertex[1],_vertex[2]);
  //_g4tracks->setVertexError(0.0,0.0,0.0);
 
  // at this point we should already have an initial pt and pz guess...
  // need to translate this into the PHG4Track object...

  vector<SimpleHit3D> track_hits;
  int clusterID;
  int clusterLayer;
  float cluster_x;
  float cluster_y;
  float cluster_z;
  //  float dEdx1;
  //  float dEdx2;

  for(unsigned int itrack=0; itrack<_tracks.size();itrack++)
  {
    SvtxTrack_v1 track;
    track.set_id(itrack);
    track_hits.clear();
    track_hits = _tracks.at(itrack).hits;
    
    for(unsigned int ihit = 0; ihit<track_hits.size();ihit++)
    {
      //      dEdx1=0;
      //      dEdx2=0;
      if( (track_hits.at(ihit).index) >= _g4clusters->size()){continue;}
      SvtxCluster *cluster = _g4clusters->get(track_hits.at(ihit).index);
      clusterID = cluster->get_id();
      clusterLayer = cluster->get_layer();
      cluster_x = cluster->get_x();
      cluster_y = cluster->get_y();
      cluster_z = cluster->get_z();
      if( (clusterLayer < (int)_seed_layers) && (clusterLayer >= 0) )
      {
        track.insert_cluster(clusterID);
      }
    }
    float kappa = _tracks.at(itrack).kappa;
    float d = _tracks.at(itrack).d;
    float phi = _tracks.at(itrack).phi;

    float dzdl = _tracks.at(itrack).dzdl;
    float z0 = _tracks.at(itrack).z0;
    
    // track.set_helix_phi(phi);
    // track.set_helix_kappa(kappa);
    // track.set_helix_d(d);
    // track.set_helix_z0(z0);
    // track.set_helix_dzdl(dzdl);
    
    float pT = kappaToPt(kappa);

    float x_center = cos(phi)*(d+1/kappa); // x coordinate of circle center
    float y_center = sin(phi)*(d+1/kappa); // y    "      "     "      "

    // find helicity from cross product sign
    short int helicity;
    if((track_hits[0].x-x_center)*(track_hits[track_hits.size()-1].y-y_center) -
       (track_hits[0].y-y_center)*(track_hits[track_hits.size()-1].x-x_center) > 0)
    {
      helicity = 1;
    }
    else
    { 
      helicity = -1;
    }
    float pZ = 0;
    if(dzdl != 1)
    {
      pZ = pT * dzdl / sqrt(1.0 - dzdl*dzdl);
    }
    int ndf = 2*_tracks.at(itrack).hits.size() - 5;
    track.set_chisq(chi_squareds[itrack]);
    track.set_ndf(ndf);
    track.set_px( pT*cos(phi-helicity*M_PI/2) );
    track.set_py( pT*sin(phi-helicity*M_PI/2) );
    track.set_pz( pZ );

    track.set_dca2d( d );
    track.set_dca2d_error(sqrt(_tracker->getKalmanStates()[itrack].C(1,1)));  

    if(_write_reco_tree==true)
    {
      _recoevent->tracks.push_back( SimpleRecoTrack() );
      _recoevent->tracks.back().px = pT*cos(phi-helicity*M_PI/2);
      _recoevent->tracks.back().py = pT*sin(phi-helicity*M_PI/2);
      _recoevent->tracks.back().pz = pZ;
      _recoevent->tracks.back().d = d;
      _recoevent->tracks.back().z0 = z0;
      _recoevent->tracks.back().quality = chi_squareds[itrack]/((float)ndf);
      _recoevent->tracks.back().charge = (-1*helicity);
    }
    

    if(_magField > 0)
    {
      track.set_charge( helicity );
    }
    else
    {
      track.set_charge( -1.0*helicity );
    }

    Matrix<float,6,6> euclidean_cov = Matrix<float,6,6>::Zero(6,6);
    convertHelixCovarianceToEuclideanCovariance( _magField, phi, d, kappa, z0, dzdl, _tracker->getKalmanStates()[itrack].C, euclidean_cov );
    
    for(unsigned int row=0;row<6;++row)
    {
      for(unsigned int col=0;col<6;++col)
      {
	track.set_error(row,col,euclidean_cov(row,col));
      }
    }

    track.set_x( vertex.get_x() + d*cos(phi) );
    track.set_y( vertex.get_y() + d*sin(phi) );
    track.set_z( vertex.get_z() + z0 );
    
    _g4tracks->insert(&track);
    vertex.insert_track(track.get_id());

    if (verbosity > 5) {
      cout << "track " << itrack << " quality = "
           << track.get_quality() << endl;
      cout << "px = " << track.get_px()
           << " py = " << track.get_py()
           << " pz = " << track.get_pz() << endl;
    }
  } // track loop

  SvtxVertex *vtxptr = _g4vertexes->insert(&vertex);
  if (verbosity > 5) vtxptr->identify();
  
  if(verbosity > 0)
  {
    cout << "PHG4HoughTransformTPC::process_event -- leaving process_event" << endl;
  }

  if(_write_reco_tree==true){ _reco_tree->Fill(); }

  _timer.get()->stop();
  return Fun4AllReturnCodes::EVENT_OK;
}
Ejemplo n.º 2
0
int PHG4HoughTransform::export_output() {

  if (_tracks.empty()) return Fun4AllReturnCodes::EVENT_OK;
  
  SvtxVertex_v1 vertex;
  vertex.set_t0(0.0);
  for (int i=0;i<3;++i) vertex.set_position(i,_vertex[i]);
  vertex.set_chisq(0.0);
  vertex.set_ndof(0); 
  vertex.set_error(0,0,0.0);
  vertex.set_error(0,1,0.0);
  vertex.set_error(0,2,0.0);
  vertex.set_error(1,0,0.0);
  vertex.set_error(1,1,0.0);
  vertex.set_error(1,2,0.0);
  vertex.set_error(2,0,0.0);
  vertex.set_error(2,1,0.0);
  vertex.set_error(2,2,0.0);

  // at this point we should already have an initial pt and pz guess...
  // need to translate this into the PHG4Track object...

  vector<SimpleHit3D> track_hits;
  int clusterID;
  int clusterLayer;

  for (unsigned int itrack = 0; itrack < _tracks.size(); itrack++) {
    SvtxTrack_v1 track;
    track.set_id(itrack);
    track_hits.clear();
    track_hits = _tracks.at(itrack).hits;

    for (unsigned int ihit = 0; ihit < track_hits.size(); ihit++) {
      if ((track_hits.at(ihit).get_id()) >= _g4clusters->size()) {
        continue;
      }
      SvtxCluster* cluster = _g4clusters->get(track_hits.at(ihit).get_id());
      clusterID = cluster->get_id();
      clusterLayer = cluster->get_layer();
      if ((clusterLayer < (int)_nlayers) && (clusterLayer >= 0)) {
        track.insert_cluster(clusterID);
      }
    }
    
    float kappa = _tracks.at(itrack).kappa;
    float d = _tracks.at(itrack).d;
    float phi = _tracks.at(itrack).phi;
    float dzdl = _tracks.at(itrack).dzdl;
    float z0 = _tracks.at(itrack).z0;

    //    track.set_helix_phi(phi);
    //    track.set_helix_kappa(kappa);
    //    track.set_helix_d(d);
    //    track.set_helix_z0(z0);
    //    track.set_helix_dzdl(dzdl);

    float pT = kappaToPt(kappa);

    float x_center =
        cos(phi) * (d + 1 / kappa);  // x coordinate of circle center
    float y_center = sin(phi) * (d + 1 / kappa);  // y    "      "     " "

    // find helicity from cross product sign
    short int helicity;
    if ((track_hits[0].get_x() - x_center) *
	(track_hits[track_hits.size() - 1].get_y() - y_center) -
	(track_hits[0].get_y() - y_center) *
	(track_hits[track_hits.size() - 1].get_x() - x_center) > 0) {
      helicity = 1;
    } else {
      helicity = -1;
    }
    float pZ = 0;
    if (dzdl != 1) {
      pZ = pT * dzdl / sqrt(1.0 - dzdl * dzdl);
    }
    int ndf = 2 * _tracks.at(itrack).hits.size() - 5;
    track.set_chisq(_track_errors[itrack]);
    track.set_ndf(ndf);
    track.set_px(pT * cos(phi - helicity * M_PI / 2));
    track.set_py(pT * sin(phi - helicity * M_PI / 2));
    track.set_pz(pZ);

    track.set_dca2d(d);
    track.set_dca2d_error(sqrt(_track_covars[itrack](1, 1)));

    if (_magField > 0) {
      track.set_charge(helicity);
    } else {
      track.set_charge(-1.0 * helicity);
    }

    Eigen::Matrix<float, 6, 6> euclidean_cov = Eigen::Matrix<float, 6, 6>::Zero(6, 6);
    convertHelixCovarianceToEuclideanCovariance(
        _magField, phi, d, kappa, z0, dzdl,
        _track_covars[itrack], euclidean_cov);

    for (unsigned int row = 0; row < 6; ++row) {
      for (unsigned int col = 0; col < 6; ++col) {
        track.set_error(row, col, euclidean_cov(row, col));
      }
    }

    track.set_x(vertex.get_x() + d * cos(phi));
    track.set_y(vertex.get_y() + d * sin(phi));
    track.set_z(vertex.get_z() + z0);

    _g4tracks->insert(&track);
    vertex.insert_track(track.get_id());

    if (verbosity > 5) {
      cout << "track " << itrack << " quality = " << track.get_quality()
           << endl;
      cout << "px = " << track.get_px() << " py = " << track.get_py()
           << " pz = " << track.get_pz() << endl;
    }
  }  // track loop

  SvtxVertex *vtxptr = _g4vertexes->insert(&vertex);
  if (verbosity > 5) vtxptr->identify();

  if (verbosity > 0) {
    cout << "PHG4HoughTransform::process_event -- leaving process_event"
         << endl;
  }

  // we are done with these now...
  _clusters.clear();
  _tracks.clear();
  _track_errors.clear();
  _track_covars.clear();
  _vertex.clear();
  _vertex.assign(3,0.0);
  
  return Fun4AllReturnCodes::EVENT_OK;
}