Index: modules/VertexFinderDA4D.h
===================================================================
--- modules/VertexFinderDA4D.h	(revision 341014c3f3bd3259d845de52bde54ba6c3b4ed08)
+++ modules/VertexFinderDA4D.h	(revision 2ad823eae91e845dd82fc75af409432061c951d9)
@@ -6,5 +6,5 @@
  *  Cluster vertices from tracks using deterministic annealing and timing information
  *
- *  \authors M. Selvaggi, L. Gray
+ *  Author O. Cerri
  *
  */
@@ -13,4 +13,5 @@
 
 #include <vector>
+#include <iostream>
 
 class TObjArray;
@@ -20,37 +21,274 @@
 class VertexFinderDA4D: public DelphesModule
 {
-public:
-  VertexFinderDA4D();
-  ~VertexFinderDA4D();
-
-  void Init();
-  void Process();
-  void Finish();
-
-  void clusterize(const TObjArray &tracks, TObjArray &clusters);
-  std::vector<Candidate *> vertices();
-
-private:
-  Bool_t fVerbose;
-  Double_t fMinPT;
-
-  Float_t fVertexSpaceSize;
-  Float_t fVertexTimeSize;
-  Bool_t fUseTc;
-  Float_t fBetaMax;
-  Float_t fBetaStop;
-  Double_t fCoolingFactor;
-  Int_t fMaxIterations;
-  Double_t fDzCutOff;
-  Double_t fD0CutOff;
-  Double_t fDtCutOff; // for when the beamspot has time
-
-  TObjArray *fInputArray;
-  TIterator *fItInputArray;
-
-  TObjArray *fOutputArray;
-  TObjArray *fVertexOutputArray;
-
-  ClassDef(VertexFinderDA4D, 1)
+  public:
+
+    class tracks_t
+    {
+      public:
+        std::vector<double> z;      // z-coordinate at point of closest approach to the beamline [mm]
+        std::vector<double> r;      // z-coordinate at point of closest approach to the beamline [mm]
+
+        std::vector<double> t;      // t-coordinate at point of closest approach to the beamline  [ps]
+
+        std::vector<double> dz2_o;    //1 over the squared error of z(pca)
+        std::vector<double> dz_o;
+        std::vector<double> dt2_o;    //1 over the squared error of t(pca)
+        std::vector<double> dt_o;
+
+        std::vector<Candidate*> tt; // a pointer to the Candidate Track
+        std::vector<double> Z;      // Normalization of P(y|x) = p_k * exp(-beta*E_ik)/Z_i
+
+        std::vector<double> w;     // track weight
+
+        // std::vector<double> pt;
+        // std::vector<double> pz;
+
+        std::vector<int> PID;
+
+        double sum_w_o_dz2 = 0;
+        double sum_w_o_dt2 = 0;
+        double sum_w = 0;
+
+        tracks_t(){}
+        ~tracks_t(){}
+
+        void addItem( double new_z, double new_t, double new_dz2_o, double new_dt2_o, Candidate * new_tt, double new_w, int new_PID)
+        {
+          z.push_back( new_z );
+          t.push_back( new_t );
+          dz2_o.push_back( new_dz2_o );
+          dz_o.push_back( sqrt(new_dz2_o) );
+          dt2_o.push_back( new_dt2_o );
+          dt_o.push_back( sqrt(new_dt2_o) );
+          tt.push_back( new_tt );
+
+          w.push_back( new_w );
+          Z.push_back( 1.0 );
+
+          PID.push_back(new_PID);
+        }
+
+        unsigned int getSize()
+        {
+          return z.size();
+        }
+
+    };
+
+    class vertex_t
+    {
+      public:
+        std::vector<double> z;  //           z coordinate
+        std::vector<double> t;  //           t coordinate
+        std::vector<double> pk; //          vertex weight for "constrained" clustering
+
+        // Elements of the covariance matrix of the posterior distribution
+        std::vector<double> szz, stt, stz;
+        std::vector<double> beta_c;
+
+        double ZSize;
+        double TSize;
+
+        // Used in the tracks-cluster assignment
+        double sum_pt;
+        std::vector<unsigned int> i_tracks;
+
+
+        vertex_t(){}
+        ~vertex_t(){}
+
+        void addItem( double new_z, double new_t, double new_pk)
+        {
+          z.push_back( new_z);
+          t.push_back( new_t);
+          pk.push_back( new_pk);
+
+          szz.push_back(0);
+          stt.push_back(0);
+          stz.push_back(0);
+          beta_c.push_back(0);
+        }
+
+        void removeItem( unsigned int i )
+        {
+          z.erase( z.begin() + i );
+          t.erase( t.begin() + i );
+          pk.erase( pk.begin() + i );
+
+          szz.erase(szz.begin() + i);
+          stt.erase(stt.begin() + i);
+          stz.erase(stz.begin() + i);
+          beta_c.erase(beta_c.begin() + i);
+        }
+
+        void mergeItems(unsigned int k1, unsigned int k2)
+        {
+          double new_pk = pk[k1] + pk[k2];
+          double new_z = (z[k1]*pk[k1] + z[k2]*pk[k2])/new_pk;
+          double new_t = (t[k1]*pk[k1] + t[k2]*pk[k2])/new_pk;
+
+          removeItem(k2);
+          z[k1] = new_z;
+          t[k1] = new_t;
+          pk[k1] = new_pk;
+        }
+
+        unsigned int getSize() const
+        {
+          return z.size();
+        }
+
+        double DistanceSquare(unsigned int k1, unsigned int k2)
+        {
+          double dz = (z[k1] - z[k2])/ZSize;
+          double dt = (t[k1] - t[k2])/TSize;
+
+          return dz*dz + dt*dt;
+        }
+
+        unsigned int NearestCluster(double t_, double z_)
+        {
+          unsigned int k_min = 0;
+          double d2_min = 0;
+
+          for(unsigned int k = 0; k < getSize(); k++)
+          {
+            double dt = (t[k] - t_)/TSize;
+            double dz = (z[k] - z_)/ZSize;
+            double d2 = dt*dt + dz*dz;
+            if(k == 0 || d2 < d2_min)
+            {
+              k_min = k;
+              d2_min = d2;
+            }
+          }
+
+          return k_min;
+        }
+
+        std::pair<double, unsigned int> ComputeAllBeta_c(unsigned int verbose = 0)
+        {
+          unsigned int nv = getSize();
+          unsigned int k_min = 0;
+          double beta_c_min = 0;
+          for(unsigned int k = 0; k < nv; k++)
+          {
+            if(szz[k] == 0 || stt[k] == 0)
+            {
+              beta_c[k] = 1000;
+              if(verbose>5)
+              {
+                std::cout << std::endl << Form("Varianca too small! beta_c set to 1000\nVertex %d:  t=%1.2e, z=%1.2e, pk=%1.2e, szz=%1.2e, stt=%1.2e, stz=%1.2e", k, t[k], z[k], pk[k], szz[k], stt[k], stz[k]);
+                // throw std::invalid_argument( "Attempting to compute beta c for uninitialized vertex" );
+              }
+            }
+
+            double aux = (stt[k] - szz[k])*(stt[k] - szz[k]) + 4*stz[k]*stz[k];
+            aux = 1. / (stt[k] + szz[k] + sqrt(aux));
+
+            if(k == 0 || aux < beta_c_min)
+            {
+              beta_c_min = aux;
+              k_min = k;
+            }
+
+            beta_c[k] = aux;
+          }
+
+          std::pair<double, unsigned int> out(beta_c_min, k_min);
+          return out;
+        }
+    };
+
+    VertexFinderDA4D();
+    ~VertexFinderDA4D();
+
+    void Init();
+    void Process();
+    void Finish();
+
+    void clusterize(TObjArray &clusters);
+
+    // Define the distance metric between tracks and vertices
+    double Energy(double t_z, double v_z, double dz2_o, double t_t, double v_t, double dt2_o);
+
+    // Fill the tracks structure from the input array
+    void fill(tracks_t &tks);
+
+    // Compute higher phase transition temperature
+    double beta0(tracks_t & tks, vertex_t &vtx);
+
+    // Compute the new vertexes position and mass
+    double update(double beta, tracks_t &tks, vertex_t &vtx, double rho0);
+
+    // If a vertex has beta_c lower than beta, split it
+    bool split(double &beta,  vertex_t & vtx, tracks_t & tks);
+
+    // Merge vertexes closer than declared dimensions
+    bool merge(vertex_t & vtx, double d2_merge);
+
+    // Eliminate clusters with only one significant/unique track
+    bool purge(vertex_t & vtx, tracks_t & tks, double & rho0, const double beta, double min_prob, double min_trk);
+
+    // Compute all the energies and set the partition function normalization for each track
+    std::vector<double> Compute_pk_exp_mBetaE(double beta, vertex_t &vtx, tracks_t &tks, double Z_init);
+
+    // Plot status of tracks and Vertices
+    void plot_status(double beta, vertex_t &vtx, tracks_t &tks, int n_it = 0, const char* flag ="");
+
+    // Plot status at the end of the fitting
+    void plot_status_end(vertex_t &vtx, tracks_t &tks);
+
+    // Plot Crush
+    void plot_split_crush(double zn, double tn, vertex_t &vtx, tracks_t &tks, int i_vtx);
+
+
+  private:
+
+    unsigned int fVerbose;
+
+    UInt_t fMaxIterations;
+    UInt_t fMaxVertexNumber;
+
+    Float_t fBetaMax;
+    Float_t fBetaStop;
+    Float_t fBetaPurge;
+
+
+    Float_t fVertexZSize;
+    Float_t fVertexTSize;
+
+    Double_t fCoolingFactor;
+
+    Double_t fDzCutOff;
+    Double_t fD0CutOff;
+    Double_t fDtCutOff;
+    Double_t fPtMin;
+    Double_t fPtMax;
+
+    Double_t fD2UpdateLim;
+    Double_t fD2Merge;
+    Double_t fMuOutlayer;
+    Double_t fMinTrackProb;
+    Int_t fMinNTrack;
+
+    TObjArray *fInputArray;
+    TIterator *fItInputArray;
+
+    TObjArray *fTrackOutputArray;
+    TObjArray *fVertexOutputArray;
+
+    ClassDef(VertexFinderDA4D, 1)
+
+    //Used to keep track of number of verices in verbose mode
+    std::vector<double> fEnergy_rec;
+    std::vector<double> fBeta_rec;
+    std::vector<double> fNvtx_rec;
+
+    // Debug purpose only
+    TObjArray *fInputGenVtx;
+    TIterator *fItInputGenVtx;
+
+    TString fFigFolderPath = "";
 };
 
