diff --git a/macro/sts/matbudget_ana.C b/macro/sts/matbudget_ana.C
index 4c9172b06bf80b589281808925b0cf9649c67860..be4851a37cd10d8e3a9336e2cb93d2feeb796aed 100644
--- a/macro/sts/matbudget_ana.C
+++ b/macro/sts/matbudget_ana.C
@@ -1,10 +1,10 @@
-/* Copyright (C) 2015-2021 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
+/* Copyright (C) 2015-2022 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
    SPDX-License-Identifier: GPL-3.0-only
-   Authors: Florian Uhlig, Volker Friese [committer], David Emschermann */
+   Authors: Eoin Clerkin, Florian Uhlig, Volker Friese [committer], David Emschermann */
 
+/*
 #if !defined(__CINT__) || defined(__MAKECINT__)
 #include "FairRadLenPoint.h"
-
 #include "TCanvas.h"
 #include "TClonesArray.h"
 #include "TFile.h"
@@ -23,6 +23,7 @@ using std::cout;
 using std::endl;
 using std::vector;
 #endif
+*/
 
 // need to set this root include path:
 // include path: -I/opt/cbm/fairsoft_jul15p1/installation/include/root -I/opt/cbm/fairroot_v-15.07-fairsoft_jul15p1/include -I/opt/cbm/fairsoft_jul15p1/installation/include
@@ -31,7 +32,7 @@ using std::vector;
 
 //Int_t matbudget_ana(Int_t nEvents = 10      , const char* stsGeo = "v16v")
 //Int_t matbudget_ana(Int_t nEvents = 1000000 , const char* stsGeo = "v16v")
-Int_t matbudget_ana(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
+Int_t matbudget_ana(Int_t nEvents = 1000000, const char* stsGeo = "v21e")
 {
 
   // Input file (MC)
@@ -54,24 +55,30 @@ Int_t matbudget_ana(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
 
   // Create output histograms
   TH1D* hisRadLen     = new TH1D("hisRadLen", "Radiation Length", 1000, 0, 100);
-  const int nStations = 8;     // number of STS stations
+
+  const int nStations =
+    8;  // 8;     // number of STS stations. I'll be adding 2 often below to account for material before the first station and after the last station.
+
   const int nBins     = 1000;  // number of bins in histograms (in both x and y)
-  const int rMax      = 55;    // maximal radius for histograms (for both x and y)
-  const double zRange = 1.4;
-  TProfile2D* hStationRadLen[nStations];
+  const int rMax      = 50;    // maximal radius for histograms (for both x and y)
+  const double zRange = 1.4;   // 1.4;
+
+  TProfile2D* hStationRadLen[nStations + 3];
   TProfile2D* hStsRadLen;
   double StsRadThick = 0;
 
   TString stsname = "Material Budget x/X_{0} [%], STS";
   hStsRadLen      = new TProfile2D(stsname, stsname, nBins, -rMax, rMax, nBins, -rMax, rMax);
 
-  for (int i = 0; i < nStations; ++i) {
+  int iStation = 0, iStationOut = 0;
+
+  for (int i = 0; i < nStations + 3; ++i) {
     TString hisname = "Radiation Thickness [%],";
     hisname += " Station";
-    hisname += i + 1;
+    hisname += i - 1;
     TString name = "Material Budget x/X_{0} [%],";
     name += " Station ";
-    name += i + 1;
+    name += i - 1;
     hStationRadLen[i] = new TProfile2D(hisname, name, nBins, -rMax, rMax, nBins, -rMax, rMax);
   }
 
@@ -79,18 +86,24 @@ Int_t matbudget_ana(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
   TVector3 vecs, veco;
   std::map<int, int> trackHitMap;
 
-
   // Event loop
   int firstEvent = 0;
   for (Int_t event = firstEvent; event < (firstEvent + nEvents) && event < tree->GetEntriesFast(); event++) {
+    // std::cout << "EVENT STARTS: " << event << std::endl;
+
     tree->GetEntry(event);
-    if (0 == event % 10000) cout << "*** Processing event " << event << endl;
+    // std::cout << "# Entries : " << tree->GetEntry(event) << std::endl;
 
+    if (0 == event % 10000) cout << "*** Processing event " << event << endl;
     const int nTracks = 1;
+    // std::cout << "# Tracks : " << nTracks << std::endl;
+
     std::vector<double> RadLengthOnTrack(nTracks, 0.0);  //trackID, vector with points on track
     std::vector<double> TrackLength(nTracks, 0.0);       //trackID, vector with points on track
 
-    vector<double> RadThick(nStations, 0);
+    // std::cout << "RadThick Reset : " << nStations+2 << std::endl;
+
+    vector<double> RadThick(nStations + 3, 0);
     Double_t x = 0;
     Double_t y = 0;
 
@@ -98,6 +111,7 @@ Int_t matbudget_ana(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
 
     // Point loop
     for (Int_t iPoint = 0; iPoint < points->GetEntriesFast(); iPoint++) {
+      // std::cout << "iPoint: " << iPoint << std::endl;
       FairRadLenPoint* point = (FairRadLenPoint*) points->At(iPoint);
 
       // Get track position at entry and exit of material
@@ -107,43 +121,63 @@ Int_t matbudget_ana(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
       posDif = posOut - posIn;
 
       // Midpoint between in and out
-      const TVector3 middle = (posOut + posIn);
-      x                     = middle.X() / 2;
-      y                     = middle.Y() / 2;
-      const double z        = middle.Z() / 2;
+      const TVector3 sum = (posOut + posIn);
+      x                  = sum.X() / 2;
+      y                  = sum.Y() / 2;
+      const double z     = sum.Z() / 2;
 
       // Material budget per unit length
       const double radThick = posDif.Mag() / point->GetRadLength();
       RadLengthOnTrack[point->GetTrackID()] += radThick;
       TrackLength[point->GetTrackID()] += posDif.Mag();
 
-      // Determine station number
-      int iStation    = posIn.Z() / 10 - 3 + 0.5;  // suppose equidistant stations at 30-100 cm
-      int iStationOut = posOut.Z() / 10 - 3 + 0.5;
-      if (iStationOut != iStation) continue;
-      if (iStation >= nStations || iStation < 0) continue;
-      RadThick[iStation] += radThick;
-    }
+      // start of the STS, start of Station 1, start of Station 2, start of Station 3, ...., start of Station 8, end of Station 8, end of STS
+      double boundaries[nStations + 3] = {-42.5, -20.5, -8.5, 2.0, 12.5, 23.0, 33.5, 44.0, 54.25, 64.75, 80.5};
+
+      iStation = 0;
+      while (posIn.Z() > boundaries[iStation] && iStation <= nStations + 1)
+        iStation++;  // iStation 2 is the first station, iStation 1 is the material before the station
+                     // inside the STS box, iStation 0 is before the STS box. Should be empty.
+
+      iStationOut = iStation;
+      while (posOut.Z() > boundaries[iStationOut] && iStationOut <= nStations + 1)
+        iStationOut++;
+
+      if (iStationOut == iStation) { RadThick[iStation] += radThick; }
+      else {
+        // Distributes part of the radThick to two stations across the boundary.
+        RadThick[iStation] += radThick * (boundaries[iStation] - posIn.Z()) / posDif.Mag();
+        int j = iStation + 1;
+        // If the material transverses more than one boundary
+        while (j < iStationOut) {
+          RadThick[j] += radThick * (boundaries[j] - boundaries[j - 1]) / posDif.Mag();
+          ++j;
+        };
+        RadThick[j] += radThick * (posOut.Z() - boundaries[j - 1]) / posDif.Mag();
+      };
+    };
 
     StsRadThick = 0;
     // Fill material budget map for each station
-    for (int i = 0; i < nStations; ++i) {
+    for (int i = 0; i <= nStations + 2; i++) {
       hStationRadLen[i]->Fill(x, y, RadThick[i] * 100);
       StsRadThick += RadThick[i];
     }
 
     hStsRadLen->Fill(x, y, StsRadThick * 100);
 
-    for (int k = 0; k < RadLengthOnTrack.size(); k++)
+    for (int k = 0; k < RadLengthOnTrack.size(); k++) {
       if (RadLengthOnTrack[k] > 0) hisRadLen->Fill(RadLengthOnTrack[k]);
+    };
 
   }  // event loop
 
+  std::cout << "EVENT LOOP ENDS " << std::endl;
 
   // Plotting the results
   //single  TCanvas* can1 = new TCanvas("c","c",800,800);
-  TCanvas* can1 = new TCanvas("c", "c", 1600, 800);
-  can1->Divide(nStations / 2, 2);
+  TCanvas* can1 = new TCanvas("c", "c", 2000, 800);
+  can1->Divide(nStations / 2 + 2, 2);
   gStyle->SetPalette(1);
   gStyle->SetOptStat(0);
 
@@ -151,13 +185,13 @@ Int_t matbudget_ana(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
   TFile* output = new TFile(outFile, "RECREATE");
   output->cd();
 
-  for (int iStation = 0; iStation < nStations; iStation++) {
+  for (int iStation = 0; iStation <= nStations + 2; iStation++) {
     can1->cd(iStation + 1);
     //single    int iStation = 7;
     hStationRadLen[iStation]->GetXaxis()->SetTitle("x [cm]");
     hStationRadLen[iStation]->GetYaxis()->SetTitle("y [cm]");
-    //hStationRadLen[iStation]->GetZaxis()->SetTitle("x/X_{0} [%]");
-    //hStationRadLen[i]->GetZaxis()->SetTitle("radiation thickness [%]");
+    hStationRadLen[iStation]->GetZaxis()->SetTitle("x/X_{0} [%]");
+    //  hStationRadLen[iStation]->GetZaxis()->SetTitle("radiation thickness [%]");
     hStationRadLen[iStation]->SetAxisRange(0, zRange, "Z");
     hStationRadLen[iStation]->Draw("colz");
     hStationRadLen[iStation]->Write();
@@ -170,11 +204,12 @@ Int_t matbudget_ana(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
   //================================================================
 
   // Plotting the results
-  TCanvas* can2 = new TCanvas("c", "c", 800, 800);
+  TCanvas* can2 = new TCanvas("c", "c", 1600, 800);
+  can2->Divide(2, 1);
   gStyle->SetPalette(1);
   gStyle->SetOptStat(0);
 
-  can2->cd();
+  //  can2->cd();
   hStsRadLen->GetXaxis()->SetTitle("x [cm]");
   hStsRadLen->GetYaxis()->SetTitle("y [cm]");
   hStsRadLen->SetAxisRange(0, 10, "Z");
@@ -187,11 +222,11 @@ Int_t matbudget_ana(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
   //================================================================
 
   TString thisStation(0);
-  can2->Clear();
-  for (int iStation = 0; iStation < nStations; iStation++) {
+  //  can2->Clear();
+  for (int iStation = 0; iStation < nStations + 3; iStation++) {
     hStationRadLen[iStation]->Draw("colz");
     // Plot file
-    thisStation.Form("%d", iStation + 1);
+    thisStation.Form("%d", iStation - 1);
     plotFile = "sts_" + stsVersion + "_station_" + thisStation + "_matbudget.png";
     can2->SaveAs(plotFile);
   }
@@ -200,6 +235,5 @@ Int_t matbudget_ana(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
   input->Close();
   output->Close();
   cout << "Material budget maps written to " << outFile << endl;
-
   return 0;
 }
diff --git a/macro/sts/matbudget_mc.C b/macro/sts/matbudget_mc.C
index 472345f9b0fa00264f520972878362ffcf05fc28..7e6b675735bdc8bc31f2eada89440cefa2db7767 100644
--- a/macro/sts/matbudget_mc.C
+++ b/macro/sts/matbudget_mc.C
@@ -1,4 +1,4 @@
-/* Copyright (C) 2015-2017 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
+/* Copyright (C) 2015-2022 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
    SPDX-License-Identifier: GPL-3.0-only
    Authors: Volker Friese [committer], David Emschermann */
 
@@ -14,16 +14,23 @@
 // T. Balog, 11.05.2015
 //
 // --------------------------------------------------------------------------
+/*
+#if !defined(__CLING__)
+#include "CbmTransport.h"
+
+#include "FairSystemInfo.h"
+
+#include "TStopwatch.h"
+#endif
+*/
+#include "../../sim/transport/gconfig/g3Config.C"
 
 //void matbudget_mc(Int_t nEvents = 10      , const char* stsGeo = "v16v")
 //void matbudget_mc(Int_t nEvents = 1000000 , const char* stsGeo = "v16v")
-void matbudget_mc(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
+void matbudget_mc(Int_t nEvents = 10000000, const char* stsGeo = "v21e")
 {
-
   // ========================================================================
   //          Adjust this part according to your requirements
-
-
   // ----- Paths and file names  --------------------------------------------
   TString stsVersion(stsGeo);
   TString inDir   = gSystem->Getenv("VMCWORKDIR");
@@ -55,13 +62,11 @@ void matbudget_mc(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
   gDebug = 0;
   // ------------------------------------------------------------------------
 
-
   // -----   Timer   --------------------------------------------------------
   TStopwatch timer;
   timer.Start();
   // ------------------------------------------------------------------------
 
-
   // -----   Create simulation run   ----------------------------------------
   FairRunSim* run = new FairRunSim();
   run->SetName("TGeant3");      // Transport engine
@@ -78,7 +83,6 @@ void matbudget_mc(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
   run->SetMaterials("media.geo");  // Materials
   // ------------------------------------------------------------------------
 
-
   // -----   Create detectors and passive volumes   -------------------------
   if (caveGeom != "") {
     FairModule* cave = new CbmCave("CAVE");
@@ -133,10 +137,8 @@ void matbudget_mc(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
     tof->SetGeometryFileName(tofGeom);
     run->AddModule(tof);
   }
-
   // ------------------------------------------------------------------------
 
-
   // -----   Create magnetic field   ----------------------------------------
   // Zero field
   CbmFieldConst* magField = new CbmFieldConst();
@@ -156,13 +158,12 @@ void matbudget_mc(Int_t nEvents = 10000000, const char* stsGeo = "v16v")
   Int_t multiplicity       = 1;  // particles per event
   FairBoxGenerator* boxGen = new FairBoxGenerator(pdgId, multiplicity);
 
-  boxGen->SetBoxXYZ(-50., -50., 50., 50., 0.);
+  boxGen->SetBoxXYZ(-50., -50., 50., 50., -200.);  // Start at -200cm in Z
   boxGen->SetPRange(0.1, 0.5);
   boxGen->SetThetaRange(0., 0.);
   boxGen->SetPhiRange(0., 360.);
 
   primGen->AddGenerator(boxGen);
-
   run->SetRadLenRegister(kTRUE);
   // ------------------------------------------------------------------------
 
diff --git a/macro/sts/matbudget_mc_phi.C b/macro/sts/matbudget_mc_phi.C
index f65af19aa3c58ab1a94308eeed362c9437860e90..5c68972daadce102c25f8ddf15d6c05f22db91c7 100644
--- a/macro/sts/matbudget_mc_phi.C
+++ b/macro/sts/matbudget_mc_phi.C
@@ -1,6 +1,6 @@
-/* Copyright (C) 2016 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
+/* Copyright (C) 2015-2022 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
    SPDX-License-Identifier: GPL-3.0-only
-   Authors: David Emschermann [committer] */
+   Authors: Thomas Balog, David Emschermann [committer] */
 
 // --------------------------------------------------------------------------
 //
@@ -15,15 +15,15 @@
 //
 // --------------------------------------------------------------------------
 
+#include "../../sim/transport/gconfig/g3Config.C"
+
 //void matbudget_mc_phi(Int_t nEvents = 10      , const char* stsGeo = "v15c")
 //void matbudget_mc_phi(Int_t nEvents = 1000000 , const char* stsGeo = "v15c")
-void matbudget_mc_phi(Int_t nEvents = 10000000, const char* stsGeo = "v15c")
+void matbudget_mc_phi(Int_t nEvents = 1000000, const char* stsGeo = "v21e")
 {
-
   // ========================================================================
   //          Adjust this part according to your requirements
 
-
   // ----- Paths and file names  --------------------------------------------
   TString stsVersion(stsGeo);
   TString inDir   = gSystem->Getenv("VMCWORKDIR");
@@ -55,13 +55,11 @@ void matbudget_mc_phi(Int_t nEvents = 10000000, const char* stsGeo = "v15c")
   gDebug = 0;
   // ------------------------------------------------------------------------
 
-
   // -----   Timer   --------------------------------------------------------
   TStopwatch timer;
   timer.Start();
   // ------------------------------------------------------------------------
 
-
   // -----   Create simulation run   ----------------------------------------
   FairRunSim* run = new FairRunSim();
   run->SetName("TGeant3");      // Transport engine
@@ -136,7 +134,6 @@ void matbudget_mc_phi(Int_t nEvents = 10000000, const char* stsGeo = "v15c")
 
   // ------------------------------------------------------------------------
 
-
   // -----   Create magnetic field   ----------------------------------------
   // Zero field
   CbmFieldConst* magField = new CbmFieldConst();
@@ -156,12 +153,9 @@ void matbudget_mc_phi(Int_t nEvents = 10000000, const char* stsGeo = "v15c")
   Int_t multiplicity       = 1;  // particles per event
   FairBoxGenerator* boxGen = new FairBoxGenerator(pdgId, multiplicity);
 
-  //  boxGen->SetBoxXYZ(-50.,-50.,50.,50.,0.);
-  boxGen->SetXYZ(0., 0., 0.);
+  boxGen->SetXYZ(0., 0., -44.0);  // The target is -44.0 cm from centre of the magnet.
   boxGen->SetPRange(0.1, 0.5);
   boxGen->SetThetaRange(0., 40.);
-  //  boxGen->SetThetaRange(2.5,25.);
-  //  boxGen->SetThetaRange(0.,0.);
   boxGen->SetPhiRange(0., 360.);
 
   primGen->AddGenerator(boxGen);
@@ -186,7 +180,6 @@ void matbudget_mc_phi(Int_t nEvents = 10000000, const char* stsGeo = "v15c")
   rtdb->print();
   // ------------------------------------------------------------------------
 
-
   // -----   Start run   ----------------------------------------------------
   run->Run(nEvents);
   // ------------------------------------------------------------------------