Update visualization.py

visualization.py

@@ -73,18 +73,51 @@ class Visualization:
         st.pyplot(fig)
         return fig
 
+    # def build_tree_diagram(self, row):
+    #     dot = Digraph()
+    #     dot.node("Q1", "Has the student attended ≥ 90% of interventions?")
+    #     dot.node("Q2", "Has the student been engaged ≥ 80% of intervention time?")
+    #     dot.node("A1", "Address Attendance", shape="box")
+    #     dot.node("A2", "Address Engagement", shape="box")
+    #     dot.node("A3", "Consider barriers, fidelity, and progress monitoring", shape="box")
+    #     if row["Attended ≥ 90%"] == "No":
+    #         dot.edge("Q1", "A1", label="No")
+    #     else:
+    #         dot.edge("Q1", "Q2", label="Yes")
+    #         dot.edge("Q2", "A2" if row["Engagement ≥ 80%"] == "No" else "A3", label="Yes")
+    #     return dot
+
     def build_tree_diagram(self, row):
         dot = Digraph()
+    
+        # Add the student's name as the first node
+        student_name = row['Student']  # Assuming 'Student' is the column name for student names
+        dot.node("Student", student_name)
+    
+        # Connect the student's name to the first question
+        dot.edge("Student", "Q1")
+    
+        # Add the first question node
         dot.node("Q1", "Has the student attended ≥ 90% of interventions?")
-        dot.node("Q2", "Has the student been engaged ≥ 80% of intervention time?")
-        dot.node("A1", "Address Attendance", shape="box")
-        dot.node("A2", "Address Engagement", shape="box")
-        dot.node("A3", "Consider barriers, fidelity, and progress monitoring", shape="box")
+    
         if row["Attended ≥ 90%"] == "No":
+            # If attendance is less than 90%, add the corresponding action node
+            dot.node("A1", "Address Attendance", shape="box")
             dot.edge("Q1", "A1", label="No")
         else:
+            # If attendance is 90% or more, proceed to the next question
+            dot.node("Q2", "Has the student been engaged ≥ 80% of intervention time?")
             dot.edge("Q1", "Q2", label="Yes")
-            dot.edge("Q2", "A2" if row["Engagement ≥ 80%"] == "No" else "A3", label="Yes")
+    
+            if row["Engagement ≥ 80%"] == "No":
+                # If engagement is less than 80%, add the corresponding action node
+                dot.node("A2", "Address Engagement", shape="box")
+                dot.edge("Q2", "A2", label="No")
+            else:
+                # If engagement is 80% or more, add the final action node
+                dot.node("A3", "Consider addressing logistical barriers", shape="box")
+                dot.edge("Q2", "A3", label="Yes")
+    
         return dot
         
     def download_chart(self, fig, filename):