From c5214093ead04f69012dd58afb285331cbc095a8 Mon Sep 17 00:00:00 2001
From: Sudipta Basak <sudiptab2100@gmail.com>
Date: Tue, 4 Jun 2024 07:02:31 +0000
Subject: [PATCH] Simulator progressVector updated

---
 DAS/observer.py  |  2 +-
 DAS/shape.py     |  1 +
 DAS/simulator.py | 36 +++++++++++++++++-------------------
 3 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/DAS/observer.py b/DAS/observer.py
index 7abc606..80b68dc 100644
--- a/DAS/observer.py
+++ b/DAS/observer.py
@@ -96,7 +96,7 @@ class Observer:
                 return np.mean(l) if l else np.NaN
 
             trafficStats = {}
-            for cl in range(0,3):
+            for cl in self.config.nodeClasses:
                 Tx = [v.statsTxInSlot for v in validators if v.nodeClass == cl]
                 Rx = [v.statsRxInSlot for v in validators if v.nodeClass == cl]
                 RxDup = [v.statsRxDupInSlot for v in validators if v.nodeClass == cl]
diff --git a/DAS/shape.py b/DAS/shape.py
index bb8ede2..dea8b7a 100644
--- a/DAS/shape.py
+++ b/DAS/shape.py
@@ -19,6 +19,7 @@ class Shape:
         self.custodyCols = custodyCols
         self.bwUplinkProd = bwUplinkProd
         self.nodeTypes = nodeTypes
+        self.nodeClasses = [0] + [_k for _k in nodeTypes.keys() if _k != "group"]
         self.randomSeed = ""
 
     def __repr__(self):
diff --git a/DAS/simulator.py b/DAS/simulator.py
index 6e7b74a..b801e77 100644
--- a/DAS/simulator.py
+++ b/DAS/simulator.py
@@ -320,12 +320,9 @@ class Simulator:
             cnN = "nodes ready"
             cnV = "validators ready"
             cnT0 = "TX builder mean"
-            cnT1 = "TX class1 mean"
-            cnT2 = "TX class2 mean"
-            cnR1 = "RX class1 mean"
-            cnR2 = "RX class2 mean"
-            cnD1 = "Dup class1 mean"
-            cnD2 = "Dup class2 mean"
+            cnT = lambda i: f"TX class{i} mean"
+            cnR = lambda i: f"RX class{i} mean"
+            cnD = lambda i: f"Dup class{i} mean"
 
             # if custody is based on the requirements of underlying individual
             # validators, we can get detailed data on how many validated.
@@ -334,19 +331,20 @@ class Simulator:
               cnVv = validatorProgress
             else:
               cnVv = validatorAllProgress
-
-            progressVector.append({
-                cnS:sampleProgress,
-                cnN:nodeProgress,
-                cnV:cnVv,
-                cnT0: trafficStats[0]["Tx"]["mean"],
-                cnT1: trafficStats[1]["Tx"]["mean"],
-                cnT2: trafficStats[2]["Tx"]["mean"],
-                cnR1: trafficStats[1]["Rx"]["mean"],
-                cnR2: trafficStats[2]["Rx"]["mean"],
-                cnD1: trafficStats[1]["RxDup"]["mean"],
-                cnD2: trafficStats[2]["RxDup"]["mean"],
-                })
+            
+            progressDict = {
+                cnS: sampleProgress,
+                cnN: nodeProgress,
+                cnV: cnVv,
+                cnT0: trafficStats[0]["Tx"]["mean"]
+            }
+            for nc in self.shape.nodeClasses:
+                if nc != 0:
+                    progressDict[cnT(nc)] = trafficStats[nc]["Tx"]["mean"]
+                    progressDict[cnR(nc)] = trafficStats[nc]["Rx"]["mean"]
+                    progressDict[cnD(nc)] = trafficStats[nc]["RxDup"]["mean"]
+            
+            progressVector.append(progressDict)
 
             if missingSamples == oldMissingSamples:
                 if len(missingVector) > self.config.steps4StopCondition: