ObjExporter.h

#include "GenFit/FitStatus.h"
#include "GenFit/GFRaveVertexFactory.h"
#include <math.h>

#include "StEvent/StEvent.h"
#include "StEvent/StFttCollection.h"
#include "StEvent/StFttCluster.h"

class StEvent;

float DEGS_TO_RAD = 3.14159f/180.0f;
float SCALE = 0.1;

class ObjExporter {
public:

    // extra output for info and debug
    const static bool verbose = false;

    // write a single vertex
    void vert( ofstream &of, float x, float y, float z ){
        of << "v " << x << " " << y << " " << z << endl;
        numVertices++;
    }
    void vert( ofstream &of, TVector3 v ){
        of << "v " << v.X() << " " << v.Y() << " " << v.Z() << endl;
        numVertices++;
    }

    // write a sphere with given sub-divsions to output file
    void sphere(TVector3 pt, float radius, int nLatitude, int nLongitude, ofstream &fout){
        int p, s, i, j;
        float x, y, z, out;
        int nPitch = nLongitude + 1;

        float pitchInc = (180. / (float)nPitch) * DEGS_TO_RAD;
        float rotInc   = (360. / (float)nLatitude) * DEGS_TO_RAD;

        //## PRINT VERTICES:
        vert(fout, pt.X(), pt.Y()+radius, pt.Z());    // Top vertex.
        vert(fout, pt.X(), pt.Y()-radius, pt.Z());    // Bottom vertex.

        int fVert = numVertices;    // Record the first vertex index for intermediate vertices.
        for(p=1; p<nPitch; p++) {    // Generate all "intermediate vertices":
            out = fabs(radius * sin((float)p * pitchInc));
            y   = radius * cos(p * pitchInc);
            for(s=0; s<nLatitude; s++) {
                x = out * cos(s * rotInc);
                z = out * sin(s * rotInc);
                fout << TString::Format("v %g %g %g\n", x+pt.X(), y+pt.Y(), z+pt.Z()).Data();
                numVertices++;
            }
        }
      
      //## OUTPUT FACES BETWEEN INTERMEDIATE POINTS:
      
      for(p=1; p<nPitch-1; p++) {
        for(s=0; s<nLatitude; s++) {
          i = p*nLatitude + s;
          j = (s==nLatitude-1) ? i-nLatitude : i;
          fout << TString::Format("f %d %d %d %d\n", 
                  (i+1-nLatitude)+fVert, (j+2-nLatitude)+fVert, (j+2)+fVert, (i+1)+fVert);
        }
      }
      
      //## OUTPUT TRIANGULAR FACES CONNECTING TO TOP AND BOTTOM VERTEX:
      
      int offLastVerts  = fVert + (nLatitude * (nLongitude-1));
      for(s=0; s<nLatitude; s++)
      {
        j = (s==nLatitude-1) ? -1 : s;
        fout << TString::Format("f %d %d %d\n", fVert-1, (j+2)+fVert,        (s+1)+fVert       ).Data();
        fout << TString::Format("f %d %d %d\n", fVert,   (s+1)+offLastVerts, (j+2)+offLastVerts).Data();
      }
    } // sphere

    // output a single face
    void face( ofstream &fout, int anchor, int v0, int v1, int v2 ){
        fout << TString::Format("f %d %d %d\n", anchor+v0, anchor+v1, anchor+v2).Data(); 
    }
    void face( ofstream &fout, int anchor, int v0, int v1, int v2, int v3 ){
        fout << TString::Format("f %d %d %d %d\n", anchor+v0, anchor+v1, anchor+v2, anchor+v3).Data(); 
    }
    // output a general prism 
    void prism( ofstream &fout, TVector3 po, TVector3 ds ){
        int fVert = numVertices;    // Record the first vertex index for intermediate vertices.
        // x-y at z=0
        vert( fout, po.X(), po.Y(), po.Z() ); // 1
        vert( fout, po.X() + ds.X(), po.Y(), po.Z() ); // 2
        vert( fout, po.X() + ds.X(), po.Y() + ds.Y(), po.Z() ); // 3
        vert( fout, po.X(), po.Y() + ds.Y(), po.Z() ); // 4

        // x-y at z=+dz
        vert( fout, po.X(), po.Y(), po.Z() + ds.Z() ); // 5
        vert( fout, po.X() + ds.X(), po.Y(), po.Z() + ds.Z() ); // 6
        vert( fout, po.X() + ds.X(), po.Y() + ds.Y(), po.Z() + ds.Z() ); // 7
        vert( fout, po.X(), po.Y() + ds.Y(), po.Z() + ds.Z() ); // 8

        face( fout, fVert, 1, 2, 3, 4 ); // bottom face 
        face( fout, fVert, 5, 6, 7, 8 ); // top face
        face( fout, fVert, 1, 2, 6, 5 ); // side 1
        face( fout, fVert, 2, 3, 7, 6 ); // side 2
        face( fout, fVert, 3, 4, 8, 7 ); // side 3
        face( fout, fVert, 1, 4, 8, 5 ); // side 4
    }

    // output a tri
    void tri( ofstream &fout, TVector3 v0, float dz, float w, float h, float phi ) {

        int fVert = numVertices;
        float dphi = 0.1/2;
        float yp = v0.Y() + (sin( phi + dphi ) + cos( phi + dphi ) ) * h;
        float yn = v0.Y() + (sin( phi - dphi ) + cos( phi - dphi ) ) * h;
        float xp = v0.X() + (cos( phi + dphi ) - sin( phi + dphi)) * h;
        float xn = v0.X() + (cos( phi - dphi ) - sin( phi - dphi)) * h;
        // top face
        vert( fout, v0.X(), v0.Y(), v0.Z() ); // 1
        vert( fout, xp, yp, v0.Z() );
        vert( fout, xn, yn, v0.Z() );

        // top face
        vert( fout, v0.X(), v0.Y(), v0.Z() - dz ); // 4
        vert( fout, xp, yp, v0.Z() - dz );
        vert( fout, xn, yn, v0.Z() - dz );

        face( fout, fVert, 1, 2, 3 ); // top
        face( fout, fVert, 4, 5, 6 ); // bottom
        face( fout, fVert, 1, 2, 5, 4 ); // side 1
        face( fout, fVert, 2, 3, 6, 5 ); // side 2
        face( fout, fVert, 3, 1, 4, 6 ); // side 3
    }

    // Compute a track projection as a linear (straight) extrapolation from two points from full projections
    static TVector3 projectAsStraightLine( genfit::Track * t, float z0, float z1, float zf, float * cov, TVector3 &mom ) {
        TVector3 tv3A = trackPosition( t, z0, cov, mom );
        TVector3 tv3B = trackPosition( t, z1, cov, mom );

        if (verbose){
            LOG_INFO << "Straight Line Projection using" << endm;
            LOG_INFO << "A.x = " << tv3A.X() << endm;
            LOG_INFO << "B.x = " << tv3B.X() << endm;
        }

        double dxdz = ( tv3B.X() - tv3A.X() ) / ( tv3B.Z() - tv3A.Z() );
        double dydz = ( tv3B.Y() - tv3A.Y() ) / ( tv3B.Z() - tv3A.Z() );

        double dx = dxdz * ( zf - z1 );
        double dy = dydz * ( zf - z1 );
        TVector3 r( tv3B.X() + dx, tv3B.Y() + dy, zf );
        return r;
    }

    // Project a track to a given z-plane and return its position, momentum, and cov matrix
    static TVector3 trackPosition( genfit::Track * t, float z, float * cov, TVector3 &mom ){

        int iPoint = 0;
        try {
            auto plane = genfit::SharedPlanePtr(
                // these normals make the planes face along z-axis
                new genfit::DetPlane(TVector3(0, 0, z), TVector3(1, 0, 0), TVector3(0, 1, 0) )
            );

            genfit::MeasuredStateOnPlane tst = t->getFittedState(iPoint);
            auto TCM = t->getCardinalRep()->get6DCov(tst);
            //  returns the track length if needed
            t->getCardinalRep()->extrapolateToPlane(tst, plane, false, true);

            TCM = t->getCardinalRep()->get6DCov(tst);

            // can get the projected positions if needed
            float x = tst.getPos().X();
            float y = tst.getPos().Y();
            float _z = tst.getPos().Z();

            mom.SetXYZ( tst.getMom().X(), tst.getMom().Y(), tst.getMom().Z() );

            if ( cov ){
                cov[0] = TCM(0,0); cov[1] = TCM(1,0); cov[2] = TCM(2,0);
                cov[3] = TCM(0,1); cov[4] = TCM(1,1); cov[5] = TCM(2,1);
                cov[6] = TCM(0,2); cov[7] = TCM(1,2); cov[8] = TCM(2,2);
            }


            return TVector3( x, y, _z );
        } catch ( genfit::Exception e ){
            LOG_INFO << "Track projection Failed from trackPoint " << iPoint  << " E: " << e.what() << endm;
            return TVector3( -990, -990, -990 );
        }
        
        
        return TVector3( -990, -990, -990 );
    }

    // Utility method
    static TVector3 trackPosition( genfit::Track * t, float z ){
        float cov[9];
        TVector3 mom;
        return trackPosition( t, z, cov, mom);
    }

    // Output the sTGC strips into a useful format for event display
    void output_ftt_strips( 
        ofstream &fout, 
        StEvent * event ){

        fout << "\n" << endl;
        fout << "o fttStrips" << endl;
        fout << "usemtl stgc_hits\n" << endl;
        float pz[] = {280.90499, 303.70498, 326.60501, 349.40499};
        TVector3 cp;
        

        for ( size_t i = 0; i < event->fttCollection()->numberOfClusters(); i++ ){
            StFttCluster* c = event->fttCollection()->clusters()[i];
            if ( c->nStrips() < 2 ) continue;
            float dw = 0.05, dlh = 60.0, dlv = 60.0;
            float mx = 0.0, my = 0.0;
            float sx = 1.0, sy = 1.0;
            
            
            if ( c->quadrant() == kFttQuadrantA ){
                mx = 0; my = 0;
                sx = 1.0; sy = 1.0;
            } else if ( c->quadrant() == kFttQuadrantB ){
                mx = 10.16*SCALE; my = 0.0*SCALE;
                sy = -1;
                dlv = -dlv;

            } else if ( c->quadrant() == kFttQuadrantC ){
                mx = -10.16*SCALE ; my = -00.0*SCALE;
                sx = -1.0; sy = -1.0;
                dlh = -dlh; dlv = -dlv;

            } else if ( c->quadrant() == kFttQuadrantD ){
                sx = -1;
                dlh = -dlh;
            }

            cp.SetZ( -pz[ c->plane() ] * SCALE );
            if ( c->orientation() == kFttHorizontal ){
                cp.SetY( my + sy * c->x()/10.0 * SCALE );
                cp.SetX( mx );
                prism( fout, cp, TVector3( dlh * SCALE, dw, dw ) );
            } else if ( c->orientation() == kFttVertical ){
                cp.SetX( mx + sx * c->x()/10.0 * SCALE );
                cp.SetY( my );
                prism( fout, cp, TVector3( dw, dlv * SCALE, dw ) );
            }
        }
    } // ftt_strips
    
    // Output the event info into a useful format for event display
    void output( std::string filename, 
            StEvent * event,
            std::vector< Seed_t> seeds, 
            std::vector< genfit::Track *> tracks, 
            const std::vector< genfit::GFRaveVertex *> &vertices, 
            std::vector<TVector3> &fttHits,
            std::vector<TVector3> &fstHits,
            std::vector<TVector3> &fcsPreHits, // EPD = preshower
            std::vector<TVector3> &fcsClusters,
            std::vector<float>    &fcsClusterEnergy
            ){
                
                LOG_INFO << "Writing: " << filename << endm;
        numVertices = 0;
        // OPEN output
        ofstream ofile( (filename + ".obj" ).c_str() );

        ofile << "\nmtllib materials.mtl\n\n" << endl;


        output_ftt_strips( ofile, event );

        // Write FWD vertices
        TVector3 startPos;
        if ( vertices.size() > 0 ){
            ofile << "o FwdVertices" << endl;
            for ( auto v : vertices ) {
                startPos.SetXYZ( v->getPos().X(), v->getPos().Y(), v->getPos().Z() );
                sphere( TVector3( v->getPos().X() * SCALE, v->getPos().Y() * SCALE, -v->getPos().Z() * SCALE ), 0.5, 10, 10, ofile );
            }
        }

        //  Write FTT hits (points)
        if ( verbose ){
            LOG_INFO << "Viz has " << fttHits.size() << " FTT Hits" << endm;
        }
        if ( fttHits.size() > 0 ){
            ofile << "\n" << endl;
            ofile << "o fttHits" << endl;
            ofile << "usemtl stgc_hits\n" << endl;
            for ( auto p : fttHits ){
                sphere( TVector3( p.X() * SCALE, p.Y() * SCALE, -p.Z() * SCALE ), 0.15, 12, 12, ofile );
            }
        }

        // write FST hits
        if ( verbose ){
            LOG_INFO << "Viz has " << fstHits.size() << " FST Hits" << endm;
        }
        if ( fstHits.size() > 0 ) {
            ofile << "\n" << endl;
            ofile << "o fstHits" << endl;
            ofile << "usemtl fst_hits\n" << endl;
            for ( auto p : fstHits ){
             
                float fstphi = TMath::ATan2( p.Y(), p.X() );
                printf( "FST PHI: %f \n", fstphi );
                // tri( ofile, TVector3( p.X() * SCALE, p.Y() * SCALE, -p.Z() * SCALE ), 0.1f, 0.1f, 3.0f, fstphi );
                sphere( TVector3( p.X() * SCALE, p.Y() * SCALE, -p.Z() * SCALE ), 0.3, 10, 10, ofile );
            }
        }

        // Output the EPD hits
        if (verbose){
            LOG_INFO << "Viz has " << fcsPreHits.size() << " EPD Hits" << endm;
        }
        if ( fcsPreHits.size() > 0 ){ 
            ofile << "\n" << endl;
            ofile << "o epd" << endl;
            ofile << "usemtl fcs_hits\n" << endl;
            for ( auto p : fcsPreHits ){
                
                sphere( TVector3( p.X() * SCALE, p.Y() * SCALE, -p.Z() * SCALE ), 0.25, 10, 10, ofile );
            }
        }

        if (verbose){
            LOG_INFO << "Viz has " << fcsClusters.size() << " FCS Hits" << endm;
        }
        if ( fcsClusters.size() > 0 ){ 
            ofile << "\n" << endl;
            ofile << "o fcs" << endl;
            ofile << "usemtl fcs_hits\n" << endl;
            int i = 0;
            for ( auto p : fcsClusters ){
                // sphere( TVector3( p.X() * SCALE, p.Y() * SCALE, -p.Z() * SCALE ), 0.75, 10, 10, ofile );
                TVector3 ds = TVector3( 1, 1, fcsClusterEnergy[i] );
                prism( ofile, TVector3( p.X() * SCALE, p.Y() * SCALE, -p.Z() * SCALE ), ds );
                i++;
            }
        }
        
        // Write the track seeds
        if (verbose){
            LOG_INFO << "Viz has " << seeds.size() << " seeds" << endm;
        }
        if ( seeds.size() > 0 ){
            ofile << "\n\no FwdSeeds\n" << endl;
            ofile << "usemtl seeds\n" << endl;
            // numVertices = 0;
            for ( auto s : seeds ) {
                size_t vStart = numVertices;
                for ( auto h : s ){

                    vert( ofile, h->getX() * SCALE, h->getY() * SCALE, -h->getZ() * SCALE );

                }

                ofile << "l ";
                for ( size_t i = vStart; i < numVertices; i++){
                    ofile << i+1 << " "; 
                }
                ofile << endl;
            }
        }

        // uncomment if you want a separate file for tracks
        // ofile.open( (filename + "_tracks.obj" ).c_str()  );
        // numVertices = 0;
        // EXPORT TRACKS

        if (verbose){
            LOG_INFO << "Viz has " << tracks.size() << " tracks Hits" << endm;
        }
        if ( tracks.size() > 0 ){
            ofile << "\n\no FwdTracks\n" << endl;
            ofile << "usemtl tracks\n" << endl;
            float zStep = 5.0; // cm
            for ( auto t : tracks ) {
                size_t vStart = numVertices;
    
                
                TVector3 lpoint;
                for ( float z = startPos.Z(); z < 875; z += zStep ){
                    TVector3 point = trackPosition( t, z );
                    if ( point.X() < -900 && point.Y() < -900 ) break;
                    if ( point.X() < -90 && point.Y() < -90 ) { z+= 50; continue;}

                    vert( ofile, point.X() * SCALE, point.Y() * SCALE, -point.Z() * SCALE );
                    lpoint = point;
                }
                
                ofile << "l ";
                for ( size_t i = vStart; i < numVertices; i++){
                    ofile << i+1 << " "; 
                }
                ofile << endl;
            } // for t in tracks
        } // if tracks.size() > 0

        ofile.close();  
    }

    size_t numVertices;

};